[2.0] make npx.load support empty .npz files (#20403)

* npx.load empty npz file * fix * set use_np * BIG_ENDIAN
apache · Jun 30, 2021 · 2a1f356 · 2a1f356
1 parent fabcd14
commit 2a1f356
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc
@@ -1974,7 +1974,8 @@ int MXNDArrayLoad(const char* fname,
         << "Failed to read 32 bits from file.";
   }
 
-  if (magic == 0x04034b50 || magic == 0x504b0304) {  // zip file format; assumed to be npz
+  if (magic == 0x04034b50 || magic == 0x504b0304 ||
+      magic == 0x06054b50 || magic == 0x504b0506) {  // zip file format; assumed to be npz
       auto[data, names] = npz::load_arrays(fname);
       ret->ret_handles.resize(data.size());
       for (size_t i = 0; i < data.size(); ++i) {

diff --git a/tests/python/unittest/test_numpy_ndarray.py b/tests/python/unittest/test_numpy_ndarray.py
@@ -1421,3 +1421,8 @@ def test_mixed_array_types_share_memory():
     mx_pinned_array = mx_array.as_in_ctx(mx.cpu_pinned(0))
     assert not _np.may_share_memory(np_array, mx_pinned_array)
     assert not _np.shares_memory(np_array, mx_pinned_array)
+
+@use_np
+def test_save_load_empty(tmp_path):
+    mx.npx.savez(str(tmp_path / 'params.npz'))
+    mx.npx.load(str(tmp_path / 'params.npz'))