raise not implemented on other subclass types

Author: drisspg

float8_experimental/float8_ops.py

@@ -10,7 +10,6 @@
 from float8_experimental.float8_python_api import addmm_float8_unwrapped
 from float8_experimental.float8_tensor import Float8Tensor
 from float8_experimental.float8_utils import is_row_major
-from torch.distributed._functional_collectives import AsyncCollectiveTensor
 from torch.utils._pytree import tree_map
 
 aten = torch.ops.aten
@@ -81,8 +80,6 @@ def preprocess_addmm(a: Float8Tensor, b: Float8Tensor):
 def float8_mm(aten_op, args, kwargs=None):
     a = args[0]
     b = args[1]
-    if isinstance(b, AsyncCollectiveTensor):
-        b = b.trigger_wait()
 
     assert isinstance(a, Float8Tensor) and isinstance(
         b, Float8Tensor

float8_experimental/float8_tensor.py

@@ -264,6 +264,15 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
         # Lazy import to avoid circular dependency
         from float8_experimental.float8_ops import FLOAT8_OPS_TABLE
 
+        # All ops in the FLOAT8_OPS_TABLE expect Float8Tensor as inputs
+        # And don't support mixed tensor subclasses. This will trigger the handler for
+        # the next type in the dispatch list. torch._C._TensorMeta is for FakeTensor
+        def allowed_subclasses(type):
+            return issubclass(cls, type) or isinstance(type, torch._C._TensorMeta)
+
+        if not all(allowed_subclasses(t) for t in types):
+            return NotImplemented
+
         if func in FLOAT8_OPS_TABLE:
             return FLOAT8_OPS_TABLE[func](func, args, kwargs)
         raise NotImplementedError(f"attempting to run {func}, this is not supported")