Add multi-dimensional arrays support for existing elementwise tests

array_api_tests/__init__.py

@@ -0,0 +1,6 @@
+from hypothesis.extra.array_api import make_strategies_namespace
+
+from . import _array_module as xp
+
+
+xps = make_strategies_namespace(xp)

array_api_tests/hypothesis_helpers.py

@@ -6,7 +6,6 @@
 from hypothesis.strategies import (lists, integers, sampled_from,
                                    shared, floats, just, composite, one_of,
                                    none, booleans)
-from hypothesis.extra.array_api import make_strategies_namespace
 
 from .pytest_helpers import nargs
 from .array_helpers import (dtype_ranges, integer_dtype_objects,
@@ -15,15 +14,12 @@
                             integer_or_boolean_dtype_objects, dtype_objects)
 from ._array_module import (full, float32, float64, bool as bool_dtype,
                             _UndefinedStub, eye, broadcast_to)
-from . import _array_module
 from . import _array_module as xp
+from . import xps
 
 from .function_stubs import elementwise_functions
 
 
-xps = make_strategies_namespace(xp)
-
-
 # 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
 # with a hypothesis health check error. Note that this functionality will not
@@ -77,10 +73,6 @@ def mutually_promotable_dtypes(draw, dtype_objects=dtype_objects):
     dtype_pairs = [(i, j) for i, j in dtype_pairs if i in dtype_objects and j in dtype_objects]
     return draw(sampled_from(dtype_pairs))
 
-shared_mutually_promotable_dtype_pairs = shared(
-    mutually_promotable_dtypes(), key="mutually_promotable_dtype_pair"
-)
-
 # shared() allows us to draw either the function or the function name and they
 # will both correspond to the same function.
 
@@ -91,10 +83,10 @@ def mutually_promotable_dtypes(draw, dtype_objects=dtype_objects):
     lambda func_name: nargs(func_name) > 1)
 
 elementwise_function_objects = elementwise_functions_names.map(
-    lambda i: getattr(_array_module, i))
+    lambda i: getattr(xp, i))
 array_functions = elementwise_function_objects
 multiarg_array_functions = multiarg_array_functions_names.map(
-    lambda i: getattr(_array_module, i))
+    lambda i: getattr(xp, i))
 
 # Limit the total size of an array shape
 MAX_ARRAY_SIZE = 10000
@@ -164,7 +156,6 @@ def two_broadcastable_shapes(draw):
 sizes = integers(0, MAX_ARRAY_SIZE)
 sqrt_sizes = integers(0, SQRT_MAX_ARRAY_SIZE)
 
-# TODO: Generate general arrays here, rather than just scalars.
 numeric_arrays = xps.arrays(
     dtype=shared(xps.floating_dtypes(), key='dtypes'),
     shape=shared(xps.array_shapes(), key='shapes'),
@@ -275,14 +266,18 @@ def multiaxis_indices(draw, shapes):
     return tuple(res)
 
 
-shared_arrays1 = xps.arrays(
-    dtype=shared_mutually_promotable_dtype_pairs.map(lambda pair: pair[0]),
-    shape=shared(two_mutually_broadcastable_shapes, key="shape_pair").map(lambda pair: pair[0]),
-)
-shared_arrays2 = xps.arrays(
-    dtype=shared_mutually_promotable_dtype_pairs.map(lambda pair: pair[1]),
-    shape=shared(two_mutually_broadcastable_shapes, key="shape_pair").map(lambda pair: pair[1]),
-)
+def two_mutual_arrays(dtypes=dtype_objects):
+    mutual_dtypes = shared(mutually_promotable_dtypes(dtypes))
+    mutual_shapes = shared(two_mutually_broadcastable_shapes)
+    arrays1 = xps.arrays(
+        dtype=mutual_dtypes.map(lambda pair: pair[0]),
+        shape=mutual_shapes.map(lambda pair: pair[0]),
+    )
+    arrays2 = xps.arrays(
+        dtype=mutual_dtypes.map(lambda pair: pair[1]),
+        shape=mutual_shapes.map(lambda pair: pair[1]),
+    )
+    return arrays1, arrays2
 
 
 @composite

array_api_tests/meta_tests/test_hypothesis_helpers.py

@@ -1,12 +1,14 @@
 from math import prod
 
 import pytest
-from hypothesis import given, strategies as st
+from hypothesis import given, strategies as st, settings
 
 from .. import _array_module as xp
 from .._array_module import _UndefinedStub
 from .. import array_helpers as ah
 from .. import hypothesis_helpers as hh
+from ..test_broadcasting import broadcast_shapes
+from ..test_elementwise_functions import sanity_check
 
 UNDEFINED_DTYPES = any(isinstance(d, _UndefinedStub) for d in ah.dtype_objects)
 pytestmark = [pytest.mark.skipif(UNDEFINED_DTYPES, reason="undefined dtypes")]
@@ -44,20 +46,24 @@ 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]
 
-    from ..test_broadcasting import broadcast_shapes
 
-    assert broadcast_shapes(pair[0], pair[1]) == pair[0]
+@given(*hh.two_mutual_arrays())
+def test_two_mutual_arrays(x1, x2):
+    sanity_check(x1, x2)
+    assert broadcast_shapes(x1.shape, x2.shape) in (x1.shape, x2.shape)
 
 
 def test_kwargs():
     results = []
 
     @given(hh.kwargs(n=st.integers(0, 10), c=st.from_regex("[a-f]")))
+    @settings(max_examples=100)
     def run(kw):
         results.append(kw)
-
     run()
+
     assert all(isinstance(kw, dict) for kw in results)
     for size in [0, 1, 2]:
         assert any(len(kw) == size for kw in results)

array_api_tests/test_creation_functions.py

@@ -8,7 +8,8 @@
                             assert_exactly_equal, isintegral, is_float_dtype)
 from .hypothesis_helpers import (numeric_dtypes, dtypes, MAX_ARRAY_SIZE,
                                  shapes, sizes, sqrt_sizes, shared_dtypes,
-                                 scalars, xps, kwargs)
+                                 scalars, kwargs)
+from . import xps
 
 from hypothesis import assume, given
 from hypothesis.strategies import integers, floats, one_of, none, booleans, just, shared, composite