FIX numpy with an array of lists

docs/changes.rst

@@ -13,6 +13,10 @@ v0.17
 -----
 - Fix loading empty NumPy object arrays: the recorded shape is preserved even
   when a zero-length axis hides later dimensions from the nested list.
+- Fix loading multi-dimensional NumPy object arrays whose cells are sequences
+  of equal length, such as lists, tuples, or arrays: the cells used to be
+  expanded into extra axes of the loaded array, changing its shape and the
+  type of its cells. :issue:`561` by `Adrin Jalali`_.
 - Fix ``skops update`` failing when a relative output path contains a
   directory, including updates with ``--inplace``.
 - Loading a skops file now checks that each part of the file holds the kind

skops/io/_numpy.py

@@ -105,22 +105,21 @@ def _construct(self):
             return content
 
         if self.type == "json":
-            # We explicitly set the dtype to "O" since we only save object
-            # arrays in json.
+            # Object arrays are stored as the nested list returned by
+            # ``ndarray.tolist()``, along with the original shape. The array is
+            # allocated with that shape and the cells are assigned one at a
+            # time, since ``np.array(nested_list, dtype="O")`` would let NumPy
+            # infer extra axes from cells which are themselves sequences of
+            # equal length, such as lists, tuples, or arrays (gh-561), and
+            # would lose the dimensions following a zero-length axis (gh-558).
             shape = self.shape.construct()
-            tmp = [o.construct() for o in self.items]
-
-            if 0 in shape:
-                return np.empty(shape, dtype="O")
-
-            # TODO: this is a hack to get the correct shape of the array. We
-            # should find _a better way_ to do this.
-            if len(shape) == 1:
-                content = np.ndarray(shape=len(tmp), dtype="O")
-                for i, v in enumerate(tmp):
-                    content[i] = v
-            else:
-                content = np.array(tmp, dtype="O")
+            items = [o.construct() for o in self.items]
+            content = np.empty(shape, dtype="O")
+            for index in np.ndindex(*shape):
+                cell = items
+                for i in index:
+                    cell = cell[i]
+                content[index] = cell
 
             return content
 

skops/io/tests/test_persist.py

@@ -808,6 +808,29 @@ def test_numpy_object_dtype_empty_array(shape, use_file, tmp_path):
     np.testing.assert_array_equal(loaded, arr)
 
 
+@pytest.mark.parametrize("shape", [(2,), (2, 2), (2, 3, 2)])
+@pytest.mark.parametrize(
+    "cell",
+    [[1, 2], (1, 2), np.array([1, 2])],
+    ids=["list", "tuple", "ndarray"],
+)
+def test_numpy_object_dtype_sequence_cells(shape, cell):
+    # Regression test for gh-561: cells of a multi-dimensional object array
+    # which are sequences of equal length used to be expanded into extra axes
+    # of the loaded array, changing both its shape and the type of its cells.
+    arr = np.empty(shape, dtype=object)
+    for index in np.ndindex(shape):
+        arr[index] = cell
+
+    loaded = loads(dumps(arr))
+
+    assert loaded.shape == shape
+    assert loaded.dtype == arr.dtype
+    for index in np.ndindex(shape):
+        assert type(loaded[index]) is type(cell)
+        np.testing.assert_array_equal(loaded[index], cell)
+
+
 def test_metainfo():
     class MyEstimator(BaseEstimator):
         """Estimator with attributes of different supported types"""