ENH: add function to get broadcast shape from a given set of shapes. (#17535)

* ENH: add function to get broadcast shape from a given set of shapes.

Add new function numpy.broadcast_shape which takes tuples
for the shapes to be broadcast against each other.
Return the broadcasted shape as a tuple.
See #17217

* Perform array allocations of size 0 for provided shape tuples

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

* Test for int as input shape

Also update docstring to include both ints and tuples of ints as input

* Remove unnecessary array_function_dispatch

* Add missing set_module

* Add release notes. Add versionadded to docstring.

Also fix up docstring details.

* follow convention for trailing comma

Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net>

* Change name to broadcast_shapes. Also add test case, and type hint.

* follow convention

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

* Update docstring

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

* Add reference to numpy docs on broadcasting to docstring

Also move versionadded

* Fix spelling

Co-authored-by: Warren Weckesser <warren.weckesser@gmail.com>

* Add broadcast_shapes to reference docs and add See Also sections

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>
Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net>
Co-authored-by: Warren Weckesser <warren.weckesser@gmail.com>

1 files changed, 63 insertions, 2 deletions

diff --git a/numpy/lib/tests/test_stride_tricks.py b/numpy/lib/tests/test_stride_tricks.py
index 9d95eb9d0..10d7a19ab 100644
--- a/numpy/lib/tests/test_stride_tricks.py
+++ b/numpy/lib/tests/test_stride_tricks.py
@@ -5,7 +5,8 @@ from numpy.testing import (
     assert_raises_regex, assert_warns,
     )
 from numpy.lib.stride_tricks import (
-    as_strided, broadcast_arrays, _broadcast_shape, broadcast_to
+    as_strided, broadcast_arrays, _broadcast_shape, broadcast_to,
+    broadcast_shapes,
     )
 
 def assert_shapes_correct(input_shapes, expected_shape):
@@ -274,7 +275,9 @@ def test_broadcast_to_raises():
 
 
 def test_broadcast_shape():
-    # broadcast_shape is already exercized indirectly by broadcast_arrays
+    # tests internal _broadcast_shape
+    # _broadcast_shape is already exercised indirectly by broadcast_arrays
+    # _broadcast_shape is also exercised by the public broadcast_shapes function
     assert_equal(_broadcast_shape(), ())
     assert_equal(_broadcast_shape([1, 2]), (2,))
     assert_equal(_broadcast_shape(np.ones((1, 1))), (1, 1))
@@ -288,6 +291,64 @@ def test_broadcast_shape():
     assert_raises(ValueError, lambda: _broadcast_shape(*bad_args))
 
 
+def test_broadcast_shapes_succeeds():
+    # tests public broadcast_shapes
+    data = [
+        [[], ()],
+        [[()], ()],
+        [[(7,)], (7,)],
+        [[(1, 2), (2,)], (1, 2)],
+        [[(1, 1)], (1, 1)],
+        [[(1, 1), (3, 4)], (3, 4)],
+        [[(6, 7), (5, 6, 1), (7,), (5, 1, 7)], (5, 6, 7)],
+        [[(5, 6, 1)], (5, 6, 1)],
+        [[(1, 3), (3, 1)], (3, 3)],
+        [[(1, 0), (0, 0)], (0, 0)],
+        [[(0, 1), (0, 0)], (0, 0)],
+        [[(1, 0), (0, 1)], (0, 0)],
+        [[(1, 1), (0, 0)], (0, 0)],
+        [[(1, 1), (1, 0)], (1, 0)],
+        [[(1, 1), (0, 1)], (0, 1)],
+        [[(), (0,)], (0,)],
+        [[(0,), (0, 0)], (0, 0)],
+        [[(0,), (0, 1)], (0, 0)],
+        [[(1,), (0, 0)], (0, 0)],
+        [[(), (0, 0)], (0, 0)],
+        [[(1, 1), (0,)], (1, 0)],
+        [[(1,), (0, 1)], (0, 1)],
+        [[(1,), (1, 0)], (1, 0)],
+        [[(), (1, 0)], (1, 0)],
+        [[(), (0, 1)], (0, 1)],
+        [[(1,), (3,)], (3,)],
+        [[2, (3, 2)], (3, 2)],
+    ]
+    for input_shapes, target_shape in data:
+        assert_equal(broadcast_shapes(*input_shapes), target_shape)
+
+    assert_equal(broadcast_shapes(*([(1, 2)] * 32)), (1, 2))
+    assert_equal(broadcast_shapes(*([(1, 2)] * 100)), (1, 2))
+
+    # regression tests for gh-5862
+    assert_equal(broadcast_shapes(*([(2,)] * 32)), (2,))
+
+
+def test_broadcast_shapes_raises():
+    # tests public broadcast_shapes
+    data = [
+        [(3,), (4,)],
+        [(2, 3), (2,)],
+        [(3,), (3,), (4,)],
+        [(1, 3, 4), (2, 3, 3)],
+        [(1, 2), (3,1), (3,2), (10, 5)],
+        [2, (2, 3)],
+    ]
+    for input_shapes in data:
+        assert_raises(ValueError, lambda: broadcast_shapes(*input_shapes))
+
+    bad_args = [(2,)] * 32 + [(3,)] * 32
+    assert_raises(ValueError, lambda: broadcast_shapes(*bad_args))
+
+
 def test_as_strided():
     a = np.array([None])
     a_view = as_strided(a)