Merge pull request #226 from honno/jax-niceties

Don't test values that are on/near the boundary of an array's dtype

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -1,15 +1,16 @@
 import re
 import itertools
 from contextlib import contextmanager
-from functools import reduce
-from math import sqrt
+from functools import reduce, wraps
+import math
 from operator import mul
-from typing import Any, List, NamedTuple, Optional, Sequence, Tuple, Union
+import struct
+from typing import Any, List, Mapping, NamedTuple, Optional, Sequence, Tuple, Union
 
 from hypothesis import assume, reject
 from hypothesis.strategies import (SearchStrategy, booleans, composite, floats,
                                    integers, just, lists, none, one_of,
-                                   sampled_from, shared)
+                                   sampled_from, shared, builds)
 
 from . import _array_module as xp, api_version
 from . import dtype_helpers as dh
@@ -20,26 +21,60 @@
 from ._array_module import broadcast_to, eye, float32, float64, full
 from .stubs import category_to_funcs
 from .pytest_helpers import nargs
-from .typing import Array, DataType, Shape
-
-# 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
-# work for floating point dtypes as those are assumed to be defined in other
-# places in the tests.
-FILTER_UNDEFINED_DTYPES = True
-# TODO: currently we assume this to be true - we probably can remove this completely
-assert FILTER_UNDEFINED_DTYPES
-
-integer_dtypes = xps.integer_dtypes() | xps.unsigned_integer_dtypes()
-floating_dtypes = xps.floating_dtypes()
-numeric_dtypes = xps.numeric_dtypes()
-integer_or_boolean_dtypes = xps.boolean_dtypes() | integer_dtypes
-boolean_dtypes = xps.boolean_dtypes()
-dtypes = xps.scalar_dtypes()
-
-shared_dtypes = shared(dtypes, key="dtype")
-shared_floating_dtypes = shared(floating_dtypes, key="dtype")
+from .typing import Array, DataType, Scalar, Shape
+
+
+def _float32ify(n: Union[int, float]) -> float:
+    n = float(n)
+    return struct.unpack("!f", struct.pack("!f", n))[0]
+
+
+@wraps(xps.from_dtype)
+def from_dtype(dtype, **kwargs) -> SearchStrategy[Scalar]:
+    """xps.from_dtype() without the crazy large numbers."""
+    if dtype == xp.bool:
+        return xps.from_dtype(dtype, **kwargs)
+
+    if dtype in dh.complex_dtypes:
+        component_dtype = dh.dtype_components[dtype]
+    else:
+        component_dtype = dtype
+
+    min_, max_ = dh.dtype_ranges[component_dtype]
+
+    if "min_value" not in kwargs.keys() and min_ != 0:
+        assert min_ < 0  # sanity check
+        min_value = -1 * math.floor(math.sqrt(abs(min_)))
+        if component_dtype == xp.float32:
+            min_value = _float32ify(min_value)
+        kwargs["min_value"] = min_value
+    if "max_value" not in kwargs.keys():
+        assert max_ > 0  # sanity check
+        max_value = math.floor(math.sqrt(max_))
+        if component_dtype == xp.float32:
+            max_value = _float32ify(max_value)
+        kwargs["max_value"] = max_value
+
+    if dtype in dh.complex_dtypes:
+        component_strat = xps.from_dtype(dh.dtype_components[dtype], **kwargs)
+        return builds(complex, component_strat, component_strat)
+    else:
+        return xps.from_dtype(dtype, **kwargs)
+
+
+@wraps(xps.arrays)
+def arrays(dtype, *args, elements=None, **kwargs) -> SearchStrategy[Array]:
+    """xps.arrays() without the crazy large numbers."""
+    if isinstance(dtype, SearchStrategy):
+        return dtype.flatmap(lambda d: arrays(d, *args, elements=elements, **kwargs))
+
+    if elements is None:
+        elements = from_dtype(dtype)
+    elif isinstance(elements, Mapping):
+        elements = from_dtype(dtype, **elements)
+
+    return xps.arrays(dtype, *args, elements=elements, **kwargs)
+
 
 _dtype_categories = [(xp.bool,), dh.uint_dtypes, dh.int_dtypes, dh.real_float_dtypes, dh.complex_dtypes]
 _sorted_dtypes = [d for category in _dtype_categories for d in category]
@@ -62,21 +97,19 @@ def _dtypes_sorter(dtype_pair: Tuple[DataType, DataType]):
     return key
 
 _promotable_dtypes = list(dh.promotion_table.keys())
-if FILTER_UNDEFINED_DTYPES:
-    _promotable_dtypes = [
-        (d1, d2) for d1, d2 in _promotable_dtypes
-        if not isinstance(d1, _UndefinedStub) or not isinstance(d2, _UndefinedStub)
-    ]
+_promotable_dtypes = [
+    (d1, d2) for d1, d2 in _promotable_dtypes
+    if not isinstance(d1, _UndefinedStub) or not isinstance(d2, _UndefinedStub)
+]
 promotable_dtypes: List[Tuple[DataType, DataType]] = sorted(_promotable_dtypes, key=_dtypes_sorter)
 
 def mutually_promotable_dtypes(
     max_size: Optional[int] = 2,
     *,
     dtypes: Sequence[DataType] = dh.all_dtypes,
 ) -> SearchStrategy[Tuple[DataType, ...]]:
-    if FILTER_UNDEFINED_DTYPES:
-        dtypes = [d for d in dtypes if not isinstance(d, _UndefinedStub)]
-        assert len(dtypes) > 0, "all dtypes undefined"  # sanity check
+    dtypes = [d for d in dtypes if not isinstance(d, _UndefinedStub)]
+    assert len(dtypes) > 0, "all dtypes undefined"  # sanity check
     if max_size == 2:
         return sampled_from(
             [(i, j) for i, j in promotable_dtypes if i in dtypes and j in dtypes]
@@ -166,7 +199,7 @@ def all_floating_dtypes() -> SearchStrategy[DataType]:
 # Limit the total size of an array shape
 MAX_ARRAY_SIZE = 10000
 # Size to use for 2-dim arrays
-SQRT_MAX_ARRAY_SIZE = int(sqrt(MAX_ARRAY_SIZE))
+SQRT_MAX_ARRAY_SIZE = int(math.sqrt(MAX_ARRAY_SIZE))
 
 # np.prod and others have overflow and math.prod is Python 3.8+ only
 def prod(seq):
@@ -202,7 +235,7 @@ def matrix_shapes(draw, stack_shapes=shapes()):
 
 @composite
 def finite_matrices(draw, shape=matrix_shapes()):
-    return draw(xps.arrays(dtype=xps.floating_dtypes(),
+    return draw(arrays(dtype=xps.floating_dtypes(),
                            shape=shape,
                            elements=dict(allow_nan=False,
                                          allow_infinity=False)))
@@ -211,7 +244,7 @@ def finite_matrices(draw, shape=matrix_shapes()):
 # Should we set a max_value here?
 _rtol_float_kw = dict(allow_nan=False, allow_infinity=False, min_value=0)
 rtols = one_of(floats(**_rtol_float_kw),
-               xps.arrays(dtype=xps.floating_dtypes(),
+               arrays(dtype=xps.floating_dtypes(),
                           shape=rtol_shared_matrix_shapes.map(lambda shape:  shape[:-2]),
                           elements=_rtol_float_kw))
 
@@ -254,7 +287,7 @@ def symmetric_matrices(draw, dtypes=xps.floating_dtypes(), finite=True):
     if not isinstance(finite, bool):
         finite = draw(finite)
     elements = {'allow_nan': False, 'allow_infinity': False} if finite else None
-    a = draw(xps.arrays(dtype=dtype, shape=shape, elements=elements))
+    a = draw(arrays(dtype=dtype, shape=shape, elements=elements))
     upper = xp.triu(a)
     lower = xp.triu(a, k=1).mT
     return upper + lower
@@ -277,7 +310,7 @@ def invertible_matrices(draw, dtypes=xps.floating_dtypes(), stack_shapes=shapes(
     n = draw(integers(0, SQRT_MAX_ARRAY_SIZE),)
     stack_shape = draw(stack_shapes)
     shape = stack_shape + (n, n)
-    d = draw(xps.arrays(dtypes, shape=n*prod(stack_shape),
+    d = draw(arrays(dtypes, shape=n*prod(stack_shape),
                         elements=dict(allow_nan=False, allow_infinity=False)))
     # Functions that require invertible matrices may do anything when it is
     # singular, including raising an exception, so we make sure the diagonals
@@ -303,7 +336,7 @@ def two_broadcastable_shapes(draw):
 sizes = integers(0, MAX_ARRAY_SIZE)
 sqrt_sizes = integers(0, SQRT_MAX_ARRAY_SIZE)
 
-numeric_arrays = xps.arrays(
+numeric_arrays = arrays(
     dtype=shared(xps.floating_dtypes(), key='dtypes'),
     shape=shared(xps.array_shapes(), key='shapes'),
 )
@@ -348,7 +381,7 @@ def python_integer_indices(draw, sizes):
 def integer_indices(draw, sizes):
     # Return either a Python integer or a 0-D array with some integer dtype
     idx = draw(python_integer_indices(sizes))
-    dtype = draw(integer_dtypes)
+    dtype = draw(xps.integer_dtypes() | xps.unsigned_integer_dtypes())
     m, M = dh.dtype_ranges[dtype]
     if m <= idx <= M:
         return draw(one_of(just(idx),
@@ -424,16 +457,15 @@ def two_mutual_arrays(
 ) -> Tuple[SearchStrategy[Array], SearchStrategy[Array]]:
     if not isinstance(dtypes, Sequence):
         raise TypeError(f"{dtypes=} not a sequence")
-    if FILTER_UNDEFINED_DTYPES:
-        dtypes = [d for d in dtypes if not isinstance(d, _UndefinedStub)]
-        assert len(dtypes) > 0  # sanity check
+    dtypes = [d for d in dtypes if not isinstance(d, _UndefinedStub)]
+    assert len(dtypes) > 0  # sanity check
     mutual_dtypes = shared(mutually_promotable_dtypes(dtypes=dtypes))
     mutual_shapes = shared(two_shapes)
-    arrays1 = xps.arrays(
+    arrays1 = arrays(
         dtype=mutual_dtypes.map(lambda pair: pair[0]),
         shape=mutual_shapes.map(lambda pair: pair[0]),
     )
-    arrays2 = xps.arrays(
+    arrays2 = arrays(
         dtype=mutual_dtypes.map(lambda pair: pair[1]),
         shape=mutual_shapes.map(lambda pair: pair[1]),
     )

array_api_tests/meta/test_hypothesis_helpers.py

@@ -128,20 +128,20 @@ def run(n, d, data):
     assert any("d" in kw.keys() and kw["d"] is xp.float64 for kw in results)
 
 
-
 @given(finite=st.booleans(), dtype=xps.floating_dtypes(), data=st.data())
 def test_symmetric_matrices(finite, dtype, data):
-    m = data.draw(hh.symmetric_matrices(st.just(dtype), finite=finite))
+    m = data.draw(hh.symmetric_matrices(st.just(dtype), finite=finite), label="m")
     assert m.dtype == dtype
     # TODO: This part of this test should be part of the .mT test
     ah.assert_exactly_equal(m, m.mT)
 
     if finite:
         ah.assert_finite(m)
 
-@given(m=hh.positive_definite_matrices(hh.shared_floating_dtypes),
-       dtype=hh.shared_floating_dtypes)
-def test_positive_definite_matrices(m, dtype):
+
+@given(dtype=xps.floating_dtypes(), data=st.data())
+def test_positive_definite_matrices(dtype, data):
+    m = data.draw(hh.positive_definite_matrices(st.just(dtype)), label="m")
     assert m.dtype == dtype
     # TODO: Test that it actually is positive definite
 

array_api_tests/test_array_object.py

@@ -24,7 +24,7 @@ def scalar_objects(
 ) -> st.SearchStrategy[Union[Scalar, List[Scalar]]]:
     """Generates scalars or nested sequences which are valid for xp.asarray()"""
     size = math.prod(shape)
-    return st.lists(xps.from_dtype(dtype), min_size=size, max_size=size).map(
+    return st.lists(hh.from_dtype(dtype), min_size=size, max_size=size).map(
         lambda l: sh.reshape(l, shape)
     )
 
@@ -123,10 +123,10 @@ def test_setitem(shape, dtypes, data):
     key = data.draw(xps.indices(shape=shape), label="key")
     _key = normalise_key(key, shape)
     axes_indices, out_shape = get_indexed_axes_and_out_shape(_key, shape)
-    value_strat = xps.arrays(dtype=dtypes.result_dtype, shape=out_shape)
+    value_strat = hh.arrays(dtype=dtypes.result_dtype, shape=out_shape)
     if out_shape == ():
         # We can pass scalars if we're only indexing one element
-        value_strat |= xps.from_dtype(dtypes.result_dtype)
+        value_strat |= hh.from_dtype(dtypes.result_dtype)
     value = data.draw(value_strat, label="value")
 
     res = xp.asarray(x, copy=True)
@@ -157,15 +157,15 @@ def test_setitem(shape, dtypes, data):
 @pytest.mark.data_dependent_shapes
 @given(hh.shapes(), st.data())
 def test_getitem_masking(shape, data):
-    x = data.draw(xps.arrays(xps.scalar_dtypes(), shape=shape), label="x")
+    x = data.draw(hh.arrays(xps.scalar_dtypes(), shape=shape), label="x")
     mask_shapes = st.one_of(
         st.sampled_from([x.shape, ()]),
         st.lists(st.booleans(), min_size=x.ndim, max_size=x.ndim).map(
             lambda l: tuple(s if b else 0 for s, b in zip(x.shape, l))
         ),
         hh.shapes(),
     )
-    key = data.draw(xps.arrays(dtype=xp.bool, shape=mask_shapes), label="key")
+    key = data.draw(hh.arrays(dtype=xp.bool, shape=mask_shapes), label="key")
 
     if key.ndim > x.ndim or not all(
         ks in (xs, 0) for xs, ks in zip(x.shape, key.shape)
@@ -201,10 +201,10 @@ def test_getitem_masking(shape, data):
 
 @given(hh.shapes(), st.data())
 def test_setitem_masking(shape, data):
-    x = data.draw(xps.arrays(xps.scalar_dtypes(), shape=shape), label="x")
-    key = data.draw(xps.arrays(dtype=xp.bool, shape=shape), label="key")
+    x = data.draw(hh.arrays(xps.scalar_dtypes(), shape=shape), label="x")
+    key = data.draw(hh.arrays(dtype=xp.bool, shape=shape), label="key")
     value = data.draw(
-        xps.from_dtype(x.dtype) | xps.arrays(dtype=x.dtype, shape=()), label="value"
+        hh.from_dtype(x.dtype) | hh.arrays(dtype=x.dtype, shape=()), label="value"
     )
 
     res = xp.asarray(x, copy=True)
@@ -263,7 +263,7 @@ def test_scalar_casting(method_name, dtype_name, stype, data):
         dtype = getattr(_xp, dtype_name)
     except AttributeError as e:
         pytest.skip(str(e))
-    x = data.draw(xps.arrays(dtype, shape=()), label="x")
+    x = data.draw(hh.arrays(dtype, shape=()), label="x")
     method = getattr(x, method_name)
     out = method()
     assert isinstance(