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

This commit is contained in:
Awni Hannun
2026-01-13 15:48:21 -08:00
committed by GitHub
parent 8654b8281d
commit 099dcc0f4c
6 changed files with 63 additions and 38 deletions
+4 -4
View File
@@ -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):