diff --git a/mlx/io/load.cpp b/mlx/io/load.cpp index 92366be7..37680cb4 100644 --- a/mlx/io/load.cpp +++ b/mlx/io/load.cpp @@ -72,6 +72,7 @@ Dtype dtype_from_array_protocol(std::string_view t) { case 'b': { if (size == 1) return bool_; + break; } case 'i': { if (size == 1) @@ -82,6 +83,7 @@ Dtype dtype_from_array_protocol(std::string_view t) { return int32; else if (size == 8) return int64; + break; } case 'u': { if (size == 1) @@ -92,21 +94,27 @@ Dtype dtype_from_array_protocol(std::string_view t) { return uint32; else if (size == 8) return uint64; + break; } case 'f': { if (size == 2) return float16; else if (size == 4) return float32; + else if (size == 8) + return float64; + break; } case 'c': { - return complex64; + if (size == 8) + return complex64; + break; } } } throw std::invalid_argument( - "[from_str] Invalid array protocol type-string: " + std::string(t)); + "[from_str] Unsupported array protocol type-string: " + std::string(t)); } #ifdef _WIN32 diff --git a/python/tests/test_load.py b/python/tests/test_load.py index eaf7bcad..b8b4b1a1 100644 --- a/python/tests/test_load.py +++ b/python/tests/test_load.py @@ -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):