Okay

Author: drisspg

float8_experimental/float8_tensor.py

@@ -101,17 +101,17 @@ def __repr__(self):
 
     def __tensor_flatten__(self):
         ctx = {
-            "_scale": self._scale,
+            # "_scale": self._scale,
             "_orig_dtype": self._orig_dtype,
         }
         # return ("_data", "_scale"), (self._orig_dtype)
-        return ["_data"], ctx
+        return ["_data", "_scale"], ctx
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors: Dict, metadata):
-        assert len(inner_tensors) == 1
-        # return Float8Tensor(tensors["_data"], tensors["_scale"], metadatas[0])
-        return Float8Tensor(inner_tensors["_data"], metadata["_scale"], metadata["_orig_dtype"])
+        assert len(inner_tensors) == 2
+        return Float8Tensor(inner_tensors["_data"], inner_tensors["_scale"], metadata["_orig_dtype"])
+        # return Float8Tensor(inner_tensors["_data"], metadata["_scale"], metadata["_orig_dtype"])
 
     def to_original_precision(self):
         return FromFloat8ConstrFunc.apply(self)