mutually_promotable_dtypes() can generate 2-or-more dtypes

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -2,13 +2,12 @@
 from operator import mul
 from math import sqrt
 import itertools
-from typing import Tuple
+from typing import Tuple, Optional
 
 from hypothesis import assume
 from hypothesis.strategies import (lists, integers, sampled_from,
                                    shared, floats, just, composite, one_of,
-                                   none, booleans)
-from hypothesis.strategies._internal.strategies import SearchStrategy
+                                   none, booleans, SearchStrategy)
 
 from .pytest_helpers import nargs
 from .array_helpers import ndindex
@@ -77,10 +76,34 @@ def _dtypes_sorter(dtype_pair):
     ]
 
 
-def mutually_promotable_dtypes(dtype_objs=dh.all_dtypes):
-    return sampled_from(
-        [(i, j) for i, j in promotable_dtypes if i in dtype_objs and j in dtype_objs]
-    )
+def mutually_promotable_dtypes(
+    max_size: Optional[int] = 2,
+    *,
+    dtypes=dh.all_dtypes,
+) -> SearchStrategy[Tuple]:
+    if max_size == 2:
+        return sampled_from(
+            [(i, j) for i, j in promotable_dtypes if i in dtypes and j in dtypes]
+        )
+    if isinstance(max_size, int) and max_size < 2:
+        raise ValueError(f'{max_size=} should be >=2')
+    strats = []
+    category_samples = {
+        category: [d for d in dtypes if d in category] for category in _dtype_categories
+    }
+    for samples in category_samples.values():
+        if len(samples) > 0:
+            strat = lists(sampled_from(samples), min_size=2, max_size=max_size)
+            strats.append(strat)
+    if len(category_samples[dh.uint_dtypes]) > 0 and len(category_samples[dh.int_dtypes]) > 0:
+        mixed_samples = category_samples[dh.uint_dtypes] + category_samples[dh.int_dtypes]
+        strat = lists(sampled_from(mixed_samples), min_size=2, max_size=max_size)
+        if xp.uint64 in mixed_samples:
+            strat = strat.filter(
+                lambda l: not (xp.uint64 in l and any(d in dh.int_dtypes for d in l))
+            )
+    return one_of(strats).map(tuple)
+
 
 # shared() allows us to draw either the function or the function name and they
 # will both correspond to the same function.
@@ -324,9 +347,9 @@ def multiaxis_indices(draw, shapes):
 
 
 def two_mutual_arrays(
-    dtype_objs=dh.all_dtypes, two_shapes=two_mutually_broadcastable_shapes
+    dtypes=dh.all_dtypes, two_shapes=two_mutually_broadcastable_shapes
 ):
-    mutual_dtypes = shared(mutually_promotable_dtypes(dtype_objs))
+    mutual_dtypes = shared(mutually_promotable_dtypes(dtypes=dtypes))
     mutual_shapes = shared(two_shapes)
     arrays1 = xps.arrays(
         dtype=mutual_dtypes.map(lambda pair: pair[0]),

array_api_tests/meta_tests/test_hypothesis_helpers.py

@@ -14,7 +14,7 @@
 UNDEFINED_DTYPES = any(isinstance(d, _UndefinedStub) for d in dh.all_dtypes)
 pytestmark = [pytest.mark.skipif(UNDEFINED_DTYPES, reason="undefined dtypes")]
 
-@given(hh.mutually_promotable_dtypes([xp.float32, xp.float64]))
+@given(hh.mutually_promotable_dtypes(dtypes=dh.float_dtypes))
 def test_mutually_promotable_dtypes(pairs):
     assert pairs in (
         (xp.float32, xp.float32),

array_api_tests/test_elementwise_functions.py

@@ -29,11 +29,11 @@
 integer_or_boolean_scalars = hh.array_scalars(hh.integer_or_boolean_dtypes)
 boolean_scalars = hh.array_scalars(hh.boolean_dtypes)
 
-two_integer_dtypes = hh.mutually_promotable_dtypes(dh.all_int_dtypes)
-two_floating_dtypes = hh.mutually_promotable_dtypes(dh.float_dtypes)
-two_numeric_dtypes = hh.mutually_promotable_dtypes(dh.numeric_dtypes)
-two_integer_or_boolean_dtypes = hh.mutually_promotable_dtypes(dh.bool_and_all_int_dtypes)
-two_boolean_dtypes = hh.mutually_promotable_dtypes((xp.bool,))
+two_integer_dtypes = hh.mutually_promotable_dtypes(dtypes=dh.all_int_dtypes)
+two_floating_dtypes = hh.mutually_promotable_dtypes(dtypes=dh.float_dtypes)
+two_numeric_dtypes = hh.mutually_promotable_dtypes(dtypes=dh.numeric_dtypes)
+two_integer_or_boolean_dtypes = hh.mutually_promotable_dtypes(dtypes=dh.bool_and_all_int_dtypes)
+two_boolean_dtypes = hh.mutually_promotable_dtypes(dtypes=(xp.bool,))
 two_any_dtypes = hh.mutually_promotable_dtypes()
 
 @st.composite

array_api_tests/test_linalg.py

@@ -110,7 +110,7 @@ def cross_args(draw, dtype_objects=dh.numeric_dtypes):
     axis = kw.get('axis', -1)
     shape[axis] = 3
 
-    mutual_dtypes = shared(mutually_promotable_dtypes(dtype_objects))
+    mutual_dtypes = shared(mutually_promotable_dtypes(dtypes=dtype_objects))
     arrays1 = xps.arrays(
         dtype=mutual_dtypes.map(lambda pair: pair[0]),
         shape=shape,
@@ -342,7 +342,7 @@ def test_matrix_transpose(x):
     _test_stacks(_array_module.matrix_transpose, x, res=res, true_val=true_val)
 
 @given(
-    *two_mutual_arrays(dtype_objs=dh.numeric_dtypes,
+    *two_mutual_arrays(dtypes=dh.numeric_dtypes,
                        two_shapes=tuples(one_d_shapes, one_d_shapes))
 )
 def test_outer(x1, x2):

array_api_tests/test_type_promotion.py

@@ -42,23 +42,7 @@ def assert_dtype(test_case: str, result_name: str, dtype: DT, expected: DT):
     assert dtype == expected, msg
 
 
-def multi_promotable_dtypes(
-    allow_bool: bool = True,
-) -> st.SearchStrategy[Tuple[DT, ...]]:
-    strats = [
-        st.lists(st.sampled_from(dh.uint_dtypes), min_size=2),
-        st.lists(st.sampled_from(dh.int_dtypes), min_size=2),
-        st.lists(st.sampled_from(dh.float_dtypes), min_size=2),
-        st.lists(st.sampled_from(dh.all_int_dtypes), min_size=2).filter(
-            lambda l: not (xp.uint64 in l and any(d in dh.int_dtypes for d in l))
-        ),
-    ]
-    if allow_bool:
-        strats.insert(0, st.lists(st.just(xp.bool), min_size=2))
-    return st.one_of(strats).map(tuple)
-
-
-@given(multi_promotable_dtypes())
+@given(hh.mutually_promotable_dtypes(None))
 def test_result_type(dtypes):
     out = xp.result_type(*dtypes)
     assert_dtype(
@@ -67,7 +51,7 @@ def test_result_type(dtypes):
 
 
 @given(
-    dtypes=multi_promotable_dtypes(allow_bool=False),
+    dtypes=hh.mutually_promotable_dtypes(None, dtypes=dh.numeric_dtypes),
     data=st.data(),
 )
 def test_meshgrid(dtypes, data):
@@ -85,7 +69,7 @@ def test_meshgrid(dtypes, data):
 
 @given(
     shape=hh.shapes(min_dims=1),
-    dtypes=multi_promotable_dtypes(allow_bool=False),
+    dtypes=hh.mutually_promotable_dtypes(None, dtypes=dh.numeric_dtypes),
     data=st.data(),
 )
 def test_concat(shape, dtypes, data):
@@ -101,7 +85,7 @@ def test_concat(shape, dtypes, data):
 
 @given(
     shape=hh.shapes(),
-    dtypes=multi_promotable_dtypes(),
+    dtypes=hh.mutually_promotable_dtypes(None),
     data=st.data(),
 )
 def test_stack(shape, dtypes, data):