Move broadcast_shapes() to shape_helpers.py

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -16,7 +16,6 @@
 from ._array_module import _UndefinedStub
 from ._array_module import bool as bool_dtype
 from ._array_module import broadcast_to, eye, float32, float64, full
-from .algos import broadcast_shapes
 from .function_stubs import elementwise_functions
 from .pytest_helpers import nargs
 from .typing import Array, DataType, Shape
@@ -243,7 +242,7 @@ def two_broadcastable_shapes(draw):
     broadcast to shape1.
     """
     shape1, shape2 = draw(two_mutually_broadcastable_shapes)
-    assume(broadcast_shapes(shape1, shape2) == shape1)
+    assume(sh.broadcast_shapes(shape1, shape2) == shape1)
     return (shape1, shape2)
 
 sizes = integers(0, MAX_ARRAY_SIZE)

array_api_tests/meta/test_broadcasting.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from ..algos import BroadcastError, _broadcast_shapes
+from .. import shape_helpers as sh
 
 
 @pytest.mark.parametrize(
@@ -19,7 +19,7 @@
     ],
 )
 def test_broadcast_shapes(shape1, shape2, expected):
-    assert _broadcast_shapes(shape1, shape2) == expected
+    assert sh._broadcast_shapes(shape1, shape2) == expected
 
 
 @pytest.mark.parametrize(
@@ -31,5 +31,5 @@ def test_broadcast_shapes(shape1, shape2, expected):
     ],
 )
 def test_broadcast_shapes_fails_on_bad_shapes(shape1, shape2):
-    with pytest.raises(BroadcastError):
-        _broadcast_shapes(shape1, shape2)
+    with pytest.raises(sh.BroadcastError):
+        sh._broadcast_shapes(shape1, shape2)

array_api_tests/meta/test_hypothesis_helpers.py

@@ -8,9 +8,9 @@
 from .. import array_helpers as ah
 from .. import dtype_helpers as dh
 from .. import hypothesis_helpers as hh
+from .. import shape_helpers as sh
 from .. import xps
 from .._array_module import _UndefinedStub
-from ..algos import broadcast_shapes
 
 UNDEFINED_DTYPES = any(isinstance(d, _UndefinedStub) for d in dh.all_dtypes)
 pytestmark = [pytest.mark.skipif(UNDEFINED_DTYPES, reason="undefined dtypes")]
@@ -62,7 +62,7 @@ def test_two_mutually_broadcastable_shapes(pair):
 def test_two_broadcastable_shapes(pair):
     for shape in pair:
         assert valid_shape(shape)
-    assert broadcast_shapes(pair[0], pair[1]) == pair[0]
+    assert sh.broadcast_shapes(pair[0], pair[1]) == pair[0]
 
 
 @given(*hh.two_mutual_arrays())

array_api_tests/pytest_helpers.py

@@ -6,7 +6,7 @@
 from . import array_helpers as ah
 from . import dtype_helpers as dh
 from . import function_stubs
-from .algos import broadcast_shapes
+from . import shape_helpers as sh
 from .typing import Array, DataType, Scalar, ScalarType, Shape
 
 __all__ = [
@@ -159,7 +159,7 @@ def assert_result_shape(
     **kw,
 ):
     if expected is None:
-        expected = broadcast_shapes(*in_shapes)
+        expected = sh.broadcast_shapes(*in_shapes)
     f_in_shapes = " . ".join(str(s) for s in in_shapes)
     f_sig = f" {f_in_shapes} "
     if kw:

array_api_tests/shape_helpers.py

@@ -8,6 +8,7 @@
 from .typing import AtomicIndex, Index, Scalar, Shape
 
 __all__ = [
+    "broadcast_shapes",
     "normalise_axis",
     "ndindex",
     "axis_ndindex",
@@ -17,6 +18,54 @@
 ]
 
 
+class BroadcastError(ValueError):
+    """Shapes do not broadcast with eachother"""
+
+
+def _broadcast_shapes(shape1: Shape, shape2: Shape) -> Shape:
+    """Broadcasts `shape1` and `shape2`"""
+    N1 = len(shape1)
+    N2 = len(shape2)
+    N = max(N1, N2)
+    shape = [None for _ in range(N)]
+    i = N - 1
+    while i >= 0:
+        n1 = N1 - N + i
+        if N1 - N + i >= 0:
+            d1 = shape1[n1]
+        else:
+            d1 = 1
+        n2 = N2 - N + i
+        if N2 - N + i >= 0:
+            d2 = shape2[n2]
+        else:
+            d2 = 1
+
+        if d1 == 1:
+            shape[i] = d2
+        elif d2 == 1:
+            shape[i] = d1
+        elif d1 == d2:
+            shape[i] = d1
+        else:
+            raise BroadcastError()
+
+        i = i - 1
+
+    return tuple(shape)
+
+
+def broadcast_shapes(*shapes: Shape):
+    if len(shapes) == 0:
+        raise ValueError("shapes=[] must be non-empty")
+    elif len(shapes) == 1:
+        return shapes[0]
+    result = _broadcast_shapes(shapes[0], shapes[1])
+    for i in range(2, len(shapes)):
+        result = _broadcast_shapes(result, shapes[i])
+    return result
+
+
 def normalise_axis(
     axis: Optional[Union[int, Tuple[int, ...]]], ndim: int
 ) -> Tuple[int, ...]:

array_api_tests/test_data_type_functions.py

@@ -9,8 +9,8 @@
 from . import dtype_helpers as dh
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
+from . import shape_helpers as sh
 from . import xps
-from .algos import broadcast_shapes
 from .typing import DataType
 
 pytestmark = pytest.mark.ci
@@ -70,7 +70,7 @@ def test_broadcast_arrays(shapes, data):
 
     out = xp.broadcast_arrays(*arrays)
 
-    out_shape = broadcast_shapes(*shapes)
+    out_shape = sh.broadcast_shapes(*shapes)
     for i, x in enumerate(arrays):
         ph.assert_dtype(
             "broadcast_arrays", x.dtype, out[i].dtype, repr_name=f"out[{i}].dtype"
@@ -90,7 +90,7 @@ def test_broadcast_to(x, data):
     shape = data.draw(
         hh.mutually_broadcastable_shapes(1, base_shape=x.shape)
         .map(lambda S: S[0])
-        .filter(lambda s: broadcast_shapes(x.shape, s) == s),
+        .filter(lambda s: sh.broadcast_shapes(x.shape, s) == s),
         label="shape",
     )
 

array_api_tests/test_linalg.py

@@ -30,8 +30,6 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 
-from .algos import broadcast_shapes
-
 from . import _array_module
 from . import _array_module as xp
 from ._array_module import linalg
@@ -56,7 +54,7 @@ def _test_stacks(f, *args, res=None, dims=2, true_val=None, **kw):
     if res is None:
         res = f(*args, **kw)
 
-    shape = args[0].shape if len(args) == 1 else broadcast_shapes(*[x.shape
+    shape = args[0].shape if len(args) == 1 else sh.broadcast_shapes(*[x.shape
                                                                     for x in args])
     for _idx in sh.ndindex(shape[:-2]):
         idx = _idx + (slice(None),)*dims
@@ -297,7 +295,7 @@ def test_matmul(x1, x2):
         assert res.shape == x1.shape[:-1]
         _test_stacks(_array_module.matmul, x1, x2, res=res, dims=1)
     else:
-        stack_shape = broadcast_shapes(x1.shape[:-2], x2.shape[:-2])
+        stack_shape = sh.broadcast_shapes(x1.shape[:-2], x2.shape[:-2])
         assert res.shape == stack_shape + (x1.shape[-2], x2.shape[-1])
         _test_stacks(_array_module.matmul, x1, x2, res=res)
 

array_api_tests/test_operators_and_elementwise_functions.py

@@ -25,7 +25,6 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from .algos import broadcast_shapes
 from .typing import Array, DataType, Param, Scalar, Shape
 
 pytestmark = pytest.mark.ci
@@ -1223,7 +1222,7 @@ def test_logical_and(x1, x2):
     out = ah.logical_and(x1, x2)
     ph.assert_dtype("logical_and", (x1.dtype, x2.dtype), out.dtype)
     # See the comments in test_equal
-    shape = broadcast_shapes(x1.shape, x2.shape)
+    shape = sh.broadcast_shapes(x1.shape, x2.shape)
     ph.assert_shape("logical_and", out.shape, shape)
     _x1 = xp.broadcast_to(x1, shape)
     _x2 = xp.broadcast_to(x2, shape)
@@ -1245,7 +1244,7 @@ def test_logical_or(x1, x2):
     out = ah.logical_or(x1, x2)
     ph.assert_dtype("logical_or", (x1.dtype, x2.dtype), out.dtype)
     # See the comments in test_equal
-    shape = broadcast_shapes(x1.shape, x2.shape)
+    shape = sh.broadcast_shapes(x1.shape, x2.shape)
     ph.assert_shape("logical_or", out.shape, shape)
     _x1 = xp.broadcast_to(x1, shape)
     _x2 = xp.broadcast_to(x2, shape)
@@ -1258,7 +1257,7 @@ def test_logical_xor(x1, x2):
     out = xp.logical_xor(x1, x2)
     ph.assert_dtype("logical_xor", (x1.dtype, x2.dtype), out.dtype)
     # See the comments in test_equal
-    shape = broadcast_shapes(x1.shape, x2.shape)
+    shape = sh.broadcast_shapes(x1.shape, x2.shape)
     ph.assert_shape("logical_xor", out.shape, shape)
     _x1 = xp.broadcast_to(x1, shape)
     _x2 = xp.broadcast_to(x2, shape)

array_api_tests/test_searching_functions.py

@@ -8,7 +8,6 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from .algos import broadcast_shapes
 
 pytestmark = pytest.mark.ci
 
@@ -134,7 +133,7 @@ def test_where(shapes, dtypes, data):
 
     out = xp.where(cond, x1, x2)
 
-    shape = broadcast_shapes(*shapes)
+    shape = sh.broadcast_shapes(*shapes)
     ph.assert_shape("where", out.shape, shape)
     # TODO: generate indices without broadcasting arrays
     _cond = xp.broadcast_to(cond, shape)

array_api_tests/test_set_functions.py

@@ -1,8 +1,8 @@
 # TODO: disable if opted out, refactor things
 import math
-import pytest
 from collections import Counter, defaultdict
 
+import pytest
 from hypothesis import assume, given
 
 from . import _array_module as xp

array_api_tests/test_sorting_functions.py

@@ -1,7 +1,7 @@
 import math
-import pytest
 from typing import Set
 
+import pytest
 from hypothesis import given
 from hypothesis import strategies as st
 from hypothesis.control import assume

array_api_tests/test_statistical_functions.py

@@ -1,7 +1,7 @@
 import math
-import pytest
 from typing import Optional
 
+import pytest
 from hypothesis import assume, given
 from hypothesis import strategies as st
 from hypothesis.control import reject