diff --git a/docs/changes.rst b/docs/changes.rst index 2fdc1539..d518f797 100644 --- a/docs/changes.rst +++ b/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 diff --git a/skops/io/_numpy.py b/skops/io/_numpy.py index db38503b..e4df49e7 100644 --- a/skops/io/_numpy.py +++ b/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 diff --git a/skops/io/tests/test_persist.py b/skops/io/tests/test_persist.py index edb5e4ad..c10c3ad7 100644 --- a/skops/io/tests/test_persist.py +++ b/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"""