Attempt to unblock SP but needs some more thought

Author: drisspg

float8_experimental/float8_ops.py

@@ -50,7 +50,38 @@ def float8_desugar_op(aten_op, args, kwargs=None):
     )
 
 
-@implements([aten.sum.dim_IntList])
+@implements([aten.split.Tensor])
+def float8_split(aten_op, args, kwargs=None):
+    new_data_tensors = aten_op(args[0]._data, *args[1:], **kwargs)
+
+    def make_float8(data):
+        return Float8Tensor(
+            data, args[0]._scale, args[0]._orig_dtype, args[0]._mm_config
+        )
+
+    out = map(make_float8, new_data_tensors)
+    return list(out)
+
+
+# Errors cant `cat_cuda float8 e4m3fn`
+# @implements([aten.cat.default])
+# def float8_cat(aten_op, args, kwargs=None):
+#     chunked_tensors: Tuple[Float8Tensor] = args[0]
+
+#     orig_dtype = args[0][0]._orig_dtype
+#     scale = args[0][0]._scale
+#     mm_config = args[0][0]._mm_config
+#     chunk_data =[]
+#     for chunk in chunked_tensors:
+#         assert chunk._orig_dtype == orig_dtype, "Expecting all chunks to be of the same dtype"
+#         assert chunk._scale is scale, "Expecting all chunks to have thee same scale as a result of a split"
+#         assert chunk._mm_config is mm_config, "Expecting all chunks to have thee same mm config as a result of a split"
+#         chunk_data.append(chunk._data)
+#     new_data = aten_op(chunk_data, *args[1:], **kwargs)
+#     return Float8Tensor(new_data, scale, orig_dtype, mm_config)
+
+
+@implements([aten.sum.dim_IntList, aten.cat.default])
 def float8_cast_up_op(aten_op, args, kwargs=None):
     """Be careful with this function, this is a "fallback" op that
     casts the output of the op to the original precision. And performs the op.