unwrap args and kwargs

Author: danielvegamyhre

torchao/prototype/moe_training/scaled_grouped_mm.py

@@ -35,6 +35,7 @@ def _scaled_grouped_mm(
         offs (int32 torch.Tensor): The offsets to use to mark the starting index of each group along dim0 of the A tensor.
         out_dtype (Optional[torch.dtype]): The dtype of the output tensor. Currently only torch.bfloat16 is supported.
     """
+    print("SCALED_GROUPED_MM")
     return _Float8GroupedMM.apply(
         A,
         B_t,

torchao/prototype/moe_training/tensor.py

@@ -42,7 +42,9 @@ def __torch_function__(cls, func, types, args, kwargs={}):
 
     @classmethod
     def __torch_dispatch__(cls, func, types, args, kwargs={}):
+        unwrap = lambda x: x._data if isinstance(x, cls) else x
         wrap = lambda x: cls(x) if isinstance(x, torch.Tensor) else x
-        output = super().__torch_dispatch__(func, types, args, kwargs)
+        unwrapped_args, unwrapped_kwargs = tree_map(unwrap, (args, kwargs))
+        output = super().__torch_dispatch__(func, types, unwrapped_args, unwrapped_kwargs)
         wrapped_output = tree_map(wrap, output)
         return wrapped_output