Redefine custom shapes strategy using xps.array_shapes

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -6,7 +6,6 @@
                                    shared, tuples as hypotheses_tuples,
                                    floats, just, composite, one_of, none,
                                    booleans, SearchStrategy)
-from hypothesis.extra.numpy import mutually_broadcastable_shapes
 from hypothesis.extra.array_api import make_strategies_namespace
 from hypothesis import assume
 
@@ -111,29 +110,25 @@ def tuples(elements, *, min_size=0, max_size=None, unique_by=None, unique=False)
     return lists(elements, min_size=min_size, max_size=max_size,
                  unique_by=unique_by, unique=unique).map(tuple)
 
-shapes = tuples(integers(0, 10)).filter(lambda shape: prod(shape) < MAX_ARRAY_SIZE)
-
 # Use this to avoid memory errors with NumPy.
 # See https://github.com/numpy/numpy/issues/15753
-shapes = tuples(integers(0, 10)).filter(
-             lambda shape: prod([i for i in shape if i]) < MAX_ARRAY_SIZE)
+shapes = xps.array_shapes(min_dims=0, min_side=0).filter(
+    lambda shape: prod(i for i in shape if i) < MAX_ARRAY_SIZE
+)
 
-two_mutually_broadcastable_shapes = mutually_broadcastable_shapes(num_shapes=2)\
+two_mutually_broadcastable_shapes = xps.mutually_broadcastable_shapes(num_shapes=2)\
     .map(lambda S: S.input_shapes)\
-    .filter(lambda S: all(prod([i for i in shape if i]) < MAX_ARRAY_SIZE for shape in S))
+    .filter(lambda S: all(prod(i for i in shape if i) < MAX_ARRAY_SIZE for shape in S))
 
 @composite
-def two_broadcastable_shapes(draw, shapes=shapes):
+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)
-    if broadcast_shapes(shape1, shape2) != shape1:
-        assume(False)
+    shape1, shape2 =  draw(two_mutually_broadcastable_shapes)
+    assume(broadcast_shapes(shape1, shape2) == shape1)
     return (shape1, shape2)
 
 sizes = integers(0, MAX_ARRAY_SIZE)

array_api_tests/meta_tests/test_hypothesis_helpers.py

@@ -1,11 +1,16 @@
+from math import prod
+
 import pytest
 from hypothesis import given
 
 from .. import _array_module as xp
 from .._array_module import _UndefinedStub
 from ..array_helpers import dtype_objects
-from ..hypothesis_helpers import (mutually_promotable_dtype_pairs,
-                                  promotable_dtypes)
+from ..hypothesis_helpers import (MAX_ARRAY_SIZE,
+                                  mutually_promotable_dtype_pairs,
+                                  promotable_dtypes, shapes,
+                                  two_broadcastable_shapes,
+                                  two_mutually_broadcastable_shapes)
 
 UNDEFINED_DTYPES = any(isinstance(d, _UndefinedStub) for d in dtype_objects)
 pytestmark = [pytest.mark.skipif(UNDEFINED_DTYPES, reason="undefined dtypes")]
@@ -27,3 +32,31 @@ def test_mutually_promotable_dtype_pairs(pairs):
         (xp.float64, xp.float64),
     )
 
+
+def valid_shape(shape) -> bool:
+    return (
+        all(isinstance(side, int) for side in shape)
+        and all(side >= 0 for side in shape)
+        and prod(shape) < MAX_ARRAY_SIZE
+    )
+
+
+@given(shapes)
+def test_shapes(shape):
+    assert valid_shape(shape)
+
+
+@given(two_mutually_broadcastable_shapes)
+def test_two_mutually_broadcastable_shapes(pair):
+    for shape in pair:
+        assert valid_shape(shape)
+
+
+@given(two_broadcastable_shapes())
+def test_two_broadcastable_shapes(pair):
+    for shape in pair:
+        assert valid_shape(shape)
+
+    from ..test_broadcasting import broadcast_shapes
+
+    assert broadcast_shapes(pair[0], pair[1]) == pair[0]

array_api_tests/test_creation_functions.py

@@ -82,7 +82,7 @@ def test_empty(shape, dtype):
 @given(
     a=xps.arrays(
         dtype=shared_dtypes,
-        shape=xps.array_shapes(),
+        shape=shapes,
     ),
     dtype=one_of(none(), promotable_dtypes(shared_dtypes)),
 )
@@ -153,7 +153,7 @@ def test_full(shape, fill_value, dtype):
 @given(
     a=xps.arrays(
         dtype=shared_dtypes,
-        shape=xps.array_shapes(),
+        shape=shapes,
     ),
     fill_value=promotable_dtypes(shared_dtypes).flatmap(xps.from_dtype),
     dtype=one_of(none(), promotable_dtypes(shared_dtypes)),
@@ -245,7 +245,7 @@ def test_ones(shape, dtype):
 @given(
     a=xps.arrays(
         dtype=shared_dtypes,
-        shape=xps.array_shapes(),
+        shape=shapes,
     ),
     dtype=one_of(none(), promotable_dtypes(shared_dtypes)),
 )
@@ -296,7 +296,7 @@ def test_zeros(shape, dtype):
 @given(
     a=xps.arrays(
         dtype=shared_dtypes,
-        shape=xps.array_shapes(),
+        shape=shapes,
     ),
     dtype=one_of(none(), promotable_dtypes(shared_dtypes)),
 )