Test broadcastable shapes for in-place operators

Author: honno

array_api_tests/meta/test_utils.py

@@ -3,11 +3,14 @@
 from hypothesis import strategies as st
 
 from .. import _array_module as xp
-from .. import xps
 from .. import shape_helpers as sh
+from .. import xps
 from ..test_creation_functions import frange
 from ..test_manipulation_functions import roll_ndindex
-from ..test_operators_and_elementwise_functions import mock_int_dtype
+from ..test_operators_and_elementwise_functions import (
+    mock_int_dtype,
+    oneway_broadcastable_shapes,
+)
 from ..test_signatures import extension_module
 
 
@@ -115,3 +118,8 @@ def test_int_to_dtype(x, dtype):
     except OverflowError:
         reject()
     assert mock_int_dtype(x, dtype) == d
+
+
+@given(oneway_broadcastable_shapes())
+def test_oneway_broadcastable_shapes(S):
+    assert sh.broadcast_shapes(*S) == S.result_shape

array_api_tests/test_operators_and_elementwise_functions.py

@@ -30,6 +30,26 @@ def boolean_and_all_integer_dtypes() -> st.SearchStrategy[DataType]:
     return xps.boolean_dtypes() | all_integer_dtypes()
 
 
+class OnewayBroadcastableShapes(NamedTuple):
+    input_shape: Shape
+    result_shape: Shape
+
+
+@st.composite
+def oneway_broadcastable_shapes(draw) -> st.SearchStrategy[OnewayBroadcastableShapes]:
+    """Return a strategy for input shapes that broadcast to result shapes."""
+    result_shape = draw(hh.shapes(min_side=1))
+    input_shape = draw(
+        xps.broadcastable_shapes(
+            result_shape,
+            # Override defaults so bad shapes are less likely to be generated.
+            max_side=None if result_shape == () else max(result_shape),
+            max_dims=len(result_shape),
+        ).filter(lambda s: sh.broadcast_shapes(result_shape, s) == result_shape)
+    )
+    return OnewayBroadcastableShapes(input_shape, result_shape)
+
+
 def mock_int_dtype(n: int, dtype: DataType) -> int:
     """Returns equivalent of `n` that mocks `dtype` behaviour."""
     nbits = dh.dtype_nbits[dtype]
@@ -326,9 +346,15 @@ def make_param(
             )
         else:
             if func_type is FuncType.IOP:
-                shared_shapes = st.shared(hh.shapes(**shapes_kw))
-                left_strat = xps.arrays(dtype=shared_dtypes, shape=shared_shapes)
-                right_strat = xps.arrays(dtype=shared_dtypes, shape=shared_shapes)
+                shared_oneway_shapes = st.shared(oneway_broadcastable_shapes())
+                left_strat = xps.arrays(
+                    dtype=shared_dtypes,
+                    shape=shared_oneway_shapes.map(lambda S: S.result_shape),
+                )
+                right_strat = xps.arrays(
+                    dtype=shared_dtypes,
+                    shape=shared_oneway_shapes.map(lambda S: S.input_shape),
+                )
             else:
                 mutual_shapes = st.shared(
                     hh.mutually_broadcastable_shapes(2, **shapes_kw)