Merge pull request #17225 from seberg/fix-broken-up-ragged-array

BUG: Fix dimension discovery of within array ragged cases

2 files changed, 23 insertions, 1 deletions

diff --git a/numpy/core/src/multiarray/array_coercion.c b/numpy/core/src/multiarray/array_coercion.c
index ffb5bd632..3f3fd1387 100644
--- a/numpy/core/src/multiarray/array_coercion.c
+++ b/numpy/core/src/multiarray/array_coercion.c
@@ -548,7 +548,7 @@ update_shape(int curr_ndim, int *max_ndim,
             success = -1;
             if (!sequence) {
                 /* Remove dimensions that we cannot use: */
-                *max_ndim -= new_ndim + i;
+                *max_ndim -= new_ndim - i;
             }
             else {
                 assert(i == 0);
diff --git a/numpy/core/tests/test_array_coercion.py b/numpy/core/tests/test_array_coercion.py
index d18df2e9c..a6c8cc8b2 100644
--- a/numpy/core/tests/test_array_coercion.py
+++ b/numpy/core/tests/test_array_coercion.py
@@ -11,6 +11,7 @@ from itertools import product
 
 import numpy as np
 from numpy.core._rational_tests import rational
+from numpy.core._multiarray_umath import _discover_array_parameters
 
 from numpy.testing import (
     assert_array_equal, assert_warns, IS_PYPY)
@@ -478,6 +479,27 @@ class TestNested:
         with pytest.raises(ValueError):
             np.array([[], np.empty((0, 1))], dtype=object)
 
+    def test_array_of_different_depths(self):
+        # When multiple arrays (or array-likes) are included in a
+        # sequences and have different depth, we currently discover
+        # as many dimensions as they share. (see also gh-17224)
+        arr = np.zeros((3, 2))
+        mismatch_first_dim = np.zeros((1, 2))
+        mismatch_second_dim = np.zeros((3, 3))
+
+        dtype, shape = _discover_array_parameters(
+            [arr, mismatch_second_dim], dtype=np.dtype("O"))
+        assert shape == (2, 3)
+
+        dtype, shape = _discover_array_parameters(
+            [arr, mismatch_first_dim], dtype=np.dtype("O"))
+        assert shape == (2,)
+        # The second case is currently supported because the arrays
+        # can be stored as objects:
+        res = np.asarray([arr, mismatch_first_dim], dtype=np.dtype("O"))
+        assert res[0] is arr
+        assert res[1] is mismatch_first_dim
+
 
 class TestBadSequences:
     # These are tests for bad objects passed into `np.array`, in general