zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-02 07:01:28 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix saving for non-contiguous arrays (#389)

Browse Source
This commit is contained in:
Awni Hannun 2024-01-06 12:44:02 -08:00 committed by GitHub
parent 608bd43604
commit b34bf5d52b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 44 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/io/load.cpp
View File

@ -50,9 +50,14 @@ void save(std::shared_ptr<io::Writer> out_stream, array a, bool retain_graph) {
throw std::invalid_argument("[save] cannot serialize an empty array");
}
if (!a.flags().contiguous) {
if (!(a.flags().row_contiguous || a.flags().col_contiguous)) {
a = reshape(flatten(a), a.shape());
a.eval(retain_graph);
}
// Check once more in-case the above ops change
if (!(a.flags().row_contiguous || a.flags().col_contiguous)) {
throw std::invalid_argument(
"[save] cannot serialize a non-contiguous array");
"[save] can only serialize row or col contiguous arrays");
}
////////////////////////////////////////////////////////

19
mlx/io/safetensor.cpp
View File

@ -142,17 +142,26 @@ void save_safetensors(
});
size_t offset = 0;
for (auto& [key, arr] : a) {
arr.eval(retain_graph_.value_or(arr.is_tracer()));
auto retain = retain_graph_.value_or(arr.is_tracer());
arr.eval(retain);
if (arr.nbytes() == 0) {
throw std::invalid_argument(
"[save_safetensors] cannot serialize an empty array key: " + key);
}
if (!arr.flags().contiguous) {
throw std::invalid_argument(
"[save_safetensors] cannot serialize a non-contiguous array key: " +
key);
// Try to make it row contiguous
if (!arr.flags().row_contiguous) {
arr = reshape(flatten(arr), arr.shape());
arr.eval(retain);
}
// Has to be row-major now but, check one more time in case
// any of the above change in the future
if (!arr.flags().row_contiguous) {
throw std::invalid_argument(
"[save_safetensors] can only serialize row-major arrays");
}
json child;
child["dtype"] = dtype_to_safetensor_str(arr.dtype());
child["shape"] = arr.shape();

23
python/tests/test_load.py
View File

@ -178,6 +178,29 @@ class TestLoad(mlx_tests.MLXTestCase):
for k, v in load_arr_mlx_npy.items():
self.assertTrue(np.array_equal(save_arrs_npy[k], v))
def test_non_contiguous(self):
if not os.path.isdir(self.test_dir):
os.mkdir(self.test_dir)
a = mx.broadcast_to(mx.array([1, 2]), [4, 2])
save_file = os.path.join(self.test_dir, "a.npy")
mx.save(save_file, a)
aload = mx.load(save_file)
self.assertTrue(mx.array_equal(a, aload))
save_file = os.path.join(self.test_dir, "a.safetensors")
mx.save_safetensors(save_file, {"a": a})
aload = mx.load(save_file)["a"]
self.assertTrue(mx.array_equal(a, aload))
# safetensors only works with row contiguous
# make sure col contiguous is handled properly
a = mx.arange(4).reshape(2, 2).T
mx.save_safetensors(save_file, {"a": a})
aload = mx.load(save_file)["a"]
self.assertTrue(mx.array_equal(a, aload))
if __name__ == "__main__":
unittest.main()