diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/stride_tricks.py | 4 | ||||
| -rw-r--r-- | numpy/lib/tests/test_stride_tricks.py | 5 |
2 files changed, 9 insertions, 0 deletions
diff --git a/numpy/lib/stride_tricks.py b/numpy/lib/stride_tricks.py index c05d56e2f..416776ff4 100644 --- a/numpy/lib/stride_tricks.py +++ b/numpy/lib/stride_tricks.py @@ -126,6 +126,10 @@ def _broadcast_shape(*args): b = np.broadcast(*args[:32]) # unfortunately, it cannot handle 32 or more arguments directly for pos in range(32, len(args), 31): + # ironically, np.broadcast does not properly handle np.broadcast + # objects (it treats them as scalars) + # use broadcasting to avoid allocating the full array + b = broadcast_to(0, b.shape) b = np.broadcast(b, *args[pos:(pos + 31)]) return b.shape diff --git a/numpy/lib/tests/test_stride_tricks.py b/numpy/lib/tests/test_stride_tricks.py index e079e0bf4..44112c970 100644 --- a/numpy/lib/tests/test_stride_tricks.py +++ b/numpy/lib/tests/test_stride_tricks.py @@ -273,6 +273,11 @@ def test_broadcast_shape(): assert_equal(_broadcast_shape(*([np.ones((1, 2))] * 32)), (1, 2)) assert_equal(_broadcast_shape(*([np.ones((1, 2))] * 100)), (1, 2)) + # regression tests for gh-5862 + assert_equal(_broadcast_shape(*([np.ones(2)] * 32 + [1])), (2,)) + bad_args = [np.ones(2)] * 32 + [np.ones(3)] * 32 + assert_raises(ValueError, lambda: _broadcast_shape(*bad_args)) + def test_as_strided(): a = np.array([None]) |