L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix numpy dtype bug (#2960)

Browse Source
This commit is contained in:
Awni Hannun
2026-01-05 11:20:40 -08:00
committed by GitHub
parent 1df6c2a009
commit 1680b6fe38
2 changed files with 25 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_load.py
+15
View File
@@ -71,6 +71,21 @@ class TestLoad(mlx_tests.MLXTestCase):
load_arr = mx.load(Path(save_file))
self.assertTrue(mx.array_equal(load_arr, save_arr))
def test_load_npy_dtype(self):
save_file = os.path.join(self.test_dir, "mlx_path.npy")
a = np.random.randn(8).astype(np.float64)
np.save(save_file, a)
out = mx.load(save_file, stream=mx.cpu)
self.assertEqual(out.dtype, mx.float64)
self.assertTrue(np.array_equal(np.array(out), a))
a = np.random.randn(8).astype(np.float64)
b = np.random.randn(8).astype(np.float64)
c = a + 0j * b
np.save(save_file, c)
with self.assertRaises(Exception):
out = mx.load(save_file, stream=mx.cpu)
def test_save_and_load_safetensors(self):
test_file = os.path.join(self.test_dir, "test.safetensors")
with self.assertRaises(Exception):