Factor out broadcast_shapes

Author: honno

array_api_tests/algos.py

@@ -0,0 +1,53 @@
+__all__ = ["broadcast_shapes"]
+
+
+from .typing import Shape
+
+
+# We use a custom exception to differentiate from potential bugs
+class BroadcastError(ValueError):
+    pass
+
+
+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

array_api_tests/hypothesis_helpers.py

@@ -19,6 +19,7 @@
 from .function_stubs import elementwise_functions
 from .pytest_helpers import nargs
 from .typing import Array, DataType, Shape
+from .algos import broadcast_shapes
 
 # Set this to True to not fail tests just because a dtype isn't implemented.
 # If no compatible dtype is implemented for a given test, the test will fail
@@ -218,7 +219,6 @@ def two_broadcastable_shapes(draw):
     This will produce two shapes (shape1, shape2) such that shape2 can be
     broadcast to shape1.
     """
-    from .test_broadcasting import broadcast_shapes
     shape1, shape2 = draw(two_mutually_broadcastable_shapes)
     assume(broadcast_shapes(shape1, shape2) == shape1)
     return (shape1, shape2)

array_api_tests/meta/test_broadcasting.py

@@ -0,0 +1,73 @@
+"""
+https://github.com/data-apis/array-api/blob/master/spec/API_specification/broadcasting.md
+"""
+
+import pytest
+from hypothesis import assume, given
+from hypothesis import strategies as st
+
+from .. import _array_module as xp
+from .. import dtype_helpers as dh
+from .. import hypothesis_helpers as hh
+from .. import pytest_helpers as ph
+from .._array_module import _UndefinedStub
+from ..algos import BroadcastError, _broadcast_shapes
+from ..function_stubs import elementwise_functions
+
+
+@pytest.mark.parametrize(
+    "shape1, shape2, expected",
+    [
+        [(8, 1, 6, 1), (7, 1, 5), (8, 7, 6, 5)],
+        [(5, 4), (1,), (5, 4)],
+        [(5, 4), (4,), (5, 4)],
+        [(15, 3, 5), (15, 1, 5), (15, 3, 5)],
+        [(15, 3, 5), (3, 5), (15, 3, 5)],
+        [(15, 3, 5), (3, 1), (15, 3, 5)],
+    ],
+)
+def test_broadcast_shapes(shape1, shape2, expected):
+    assert _broadcast_shapes(shape1, shape2) == expected
+
+
+@pytest.mark.parametrize(
+    "shape1, shape2",
+    [
+        [(3,), (4,)],  # dimension does not match
+        [(2, 1), (8, 4, 3)],  # second dimension does not match
+        [(15, 3, 5), (15, 3)],  # singleton dimensions can only be prepended
+    ],
+)
+def test_broadcast_shapes_fails_on_bad_shapes(shape1, shape2):
+    with pytest.raises(BroadcastError):
+        _broadcast_shapes(shape1, shape2)
+
+
+# TODO: Extend this to all functions (not just elementwise), and handle
+# functions that take more than 2 args
+@pytest.mark.parametrize(
+    "func_name", [i for i in elementwise_functions.__all__ if ph.nargs(i) > 1]
+)
+@given(shape1=hh.shapes(), shape2=hh.shapes(), data=st.data())
+def test_broadcasting_hypothesis(func_name, shape1, shape2, data):
+    dtype = data.draw(st.sampled_from(dh.func_in_dtypes[func_name]), label="dtype")
+    if hh.FILTER_UNDEFINED_DTYPES:
+        assume(not isinstance(dtype, _UndefinedStub))
+    func = getattr(xp, func_name)
+    if isinstance(func, xp._UndefinedStub):
+        func._raise()
+    args = [xp.ones(shape1, dtype=dtype), xp.ones(shape2, dtype=dtype)]
+    try:
+        broadcast_shape = _broadcast_shapes(shape1, shape2)
+    except BroadcastError:
+        ph.raises(
+            Exception,
+            lambda: func(*args),
+            f"{func_name} should raise an exception from not being able to broadcast inputs with hh.shapes {(shape1, shape2)}",
+        )
+    else:
+        result = ph.doesnt_raise(
+            lambda: func(*args),
+            f"{func_name} raised an unexpected exception from broadcastable inputs with hh.shapes {(shape1, shape2)}",
+        )
+        assert result.shape == broadcast_shape, "broadcast hh.shapes incorrect"

array_api_tests/meta/test_hypothesis_helpers.py

@@ -1,15 +1,16 @@
 from math import prod
 
 import pytest
-from hypothesis import given, strategies as st, settings
+from hypothesis import given, settings
+from hypothesis import strategies as st
 
 from .. import _array_module as xp
-from .. import xps
-from .._array_module import _UndefinedStub
 from .. import array_helpers as ah
 from .. import dtype_helpers as dh
 from .. import hypothesis_helpers as hh
-from ..test_broadcasting import broadcast_shapes
+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")]

array_api_tests/pytest_helpers.py

@@ -1,3 +1,4 @@
+from array_api_tests.algos import broadcast_shapes
 import math
 from inspect import getfullargspec
 from typing import Any, Dict, Optional, Tuple, Union
@@ -17,6 +18,7 @@
     "assert_default_float",
     "assert_default_int",
     "assert_shape",
+    "assert_result_shape",
     "assert_fill",
 ]
 
@@ -131,6 +133,28 @@ def assert_shape(
     assert out_shape == expected, msg
 
 
+def assert_result_shape(
+    func_name: str,
+    in_shapes: Tuple[Shape],
+    out_shape: Shape,
+    /,
+    expected: Optional[Shape] = None,
+    *,
+    repr_name="out.shape",
+    **kw,
+):
+    if expected is None:
+        expected = broadcast_shapes(*in_shapes)
+    f_in_shapes = " . ".join(str(s) for s in in_shapes)
+    f_sig = f" {f_in_shapes} "
+    if kw:
+        f_sig += f", {fmt_kw(kw)}"
+    msg = (
+        f"{repr_name}={out_shape}, but should be {expected} " f"[{func_name}({f_sig})]"
+    )
+    assert out_shape == expected, msg
+
+
 def assert_fill(
     func_name: str, fill_value: Scalar, dtype: DataType, out: Array, /, **kw
 ):

array_api_tests/test_broadcasting.py

@@ -24,12 +24,9 @@
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
 from . import xps
+from .algos import broadcast_shapes
 from .typing import Array, DataType, Param, Scalar
 
-# We might as well use this implementation rather than xp.broadcast_shapes()
-from .test_broadcasting import broadcast_shapes
-
-
 # When appropiate, this module tests operators alongside their respective
 # elementwise methods. We do this by parametrizing a generalised test method
 # with every relevant method and operator.

array_api_tests/test_elementwise_functions.py

@@ -29,7 +29,7 @@
 from . import dtype_helpers as dh
 from . import pytest_helpers as ph
 
-from .test_broadcasting import broadcast_shapes
+from .algos import broadcast_shapes
 
 from . import _array_module
 from ._array_module import linalg