Expose to/from fp8 in Python and don't auto-convert fp8 when loading from safetensors (#2985)

python/src/ops.cpp

@@ -5479,10 +5479,10 @@ void init_ops(nb::module_& m) {
         If ``w`` is expected to receive gradients, it must be provided in
         non-quantized form.
 
-        If ``x`` and `w`` are not quantized, their data types must be ``float32``, 
+        If ``x`` and `w`` are not quantized, their data types must be ``float32``,
         ``float16``, or ``bfloat16``.
         If ``w`` is quantized, it must be packed in unsigned integers.
-        
+
       Args:
         x (array): Input array.
         w (array): Weight matrix. If quantized, it is packed in unsigned integers.
@@ -5502,4 +5502,40 @@ void init_ops(nb::module_& m) {
         array: The result of the multiplication of quantized ``x`` with quantized ``w``.
         needed).
   )pbdoc");
+  m.def(
+      "from_fp8",
+      &mx::from_fp8,
+      nb::arg(),
+      "dtype"_a = mx::bfloat16,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def from_fp8(x: array, dtype: Dtype = bfloat16, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+      Convert the array from fp8 (e4m3) to another floating-point type.
+
+      Args:
+        x (array): The input fp8 array with type ``uint8``.
+        dtype (Dtype): The data type to convert to. Default: ``bfloat16``.
+
+      Returns:
+        array: The array converted from fp8.
+  )pbdoc");
+  m.def(
+      "to_fp8",
+      &mx::to_fp8,
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def to_fp8(x: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+      Convert the array to fp8 (e4m3) from another floating-point type.
+
+      Args:
+        x (array): The input array.
+
+      Returns:
+        array: The array converted to fp8 with type ``uint8``.
+  )pbdoc");
 }

python/tests/cuda_skip.py

@@ -1,5 +1,4 @@
 cuda_skip = {
-    "TestLoad.test_load_f8_e4m3",
     "TestLayers.test_quantized_embedding",
     # Block masked matmul NYI
     "TestBlas.test_block_masked_matmul",

python/tests/test_load.py

@@ -168,8 +168,8 @@ class TestLoad(mlx_tests.MLXTestCase):
 
         expected = [
             0,
-            mx.nan,
-            mx.nan,
+            448,
+            -448,
             -0.875,
             0.4375,
             -0.005859,
@@ -179,12 +179,12 @@ class TestLoad(mlx_tests.MLXTestCase):
             -0.0039,
         ]
         expected = mx.array(expected, dtype=mx.bfloat16)
-        contents = b'H\x00\x00\x00\x00\x00\x00\x00{"tensor":{"dtype":"F8_E4M3","shape":[10],"data_offsets":[0,10]}}       \x00\x7f\xff\xb6.\x83\xba\xba\xbc\x82'
+        contents = b'H\x00\x00\x00\x00\x00\x00\x00{"tensor":{"dtype":"F8_E4M3","shape":[10],"data_offsets":[0,10]}}       \x00~\xfe\xb6.\x83\xba\xba\xbc\x82'
         with tempfile.NamedTemporaryFile(suffix=".safetensors") as f:
             f.write(contents)
             f.seek(0)
             out = mx.load(f)["tensor"]
-        self.assertTrue(mx.allclose(out[0], expected[0], equal_nan=True))
+        self.assertTrue(mx.allclose(mx.from_fp8(out), expected))
 
     def test_save_and_load_gguf_metadata_basic(self):
         if not os.path.isdir(self.test_dir):

python/tests/test_ops.py

@@ -3197,8 +3197,6 @@ class TestOps(mlx_tests.MLXTestCase):
             )
         )
 
-
-class TestBroadcast(mlx_tests.MLXTestCase):
     def test_broadcast_shapes(self):
         # Basic broadcasting
         self.assertEqual(mx.broadcast_shapes((1, 2, 3), (3,)), (1, 2, 3))
@@ -3243,6 +3241,13 @@ class TestBroadcast(mlx_tests.MLXTestCase):
         self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
         x = mx.array([3.0, mx.nan, 2.0, 0.0]) + 1j * mx.array([1.0] * 4)
 
+    def test_to_from_fp8(self):
+        vals = mx.array(
+            [448, 256, 192, 128, 96, 64, 48, 32, 24, 16, 12, 8, 6, 4, 3, 2, 0.015625]
+        )
+        self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals))
+        self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()