Prevent hypothesis testing boundary numbers

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,7 +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
+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]
@@ -145,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):
@@ -181,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)))
@@ -190,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))
 
@@ -233,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
@@ -256,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
@@ -282,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'),
 )
@@ -407,11 +461,11 @@ def two_mutual_arrays(
     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/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(

array_api_tests/test_creation_functions.py

@@ -86,8 +86,8 @@ def test_arange(dtype, data):
         start = data.draw(reals(), label="start")
         stop = data.draw(reals() | st.none(), label="stop")
     else:
-        start = data.draw(xps.from_dtype(dtype), label="start")
-        stop = data.draw(xps.from_dtype(dtype), label="stop")
+        start = data.draw(hh.from_dtype(dtype), label="start")
+        stop = data.draw(hh.from_dtype(dtype), label="stop")
     if stop is None:
         _start = 0
         _stop = start
@@ -107,9 +107,9 @@ def test_arange(dtype, data):
         step_strats = []
         if dtype in dh.int_dtypes:
             step_min = min(math.floor(-tol), -1)
-            step_strats.append(xps.from_dtype(dtype, max_value=step_min))
+            step_strats.append(hh.from_dtype(dtype, max_value=step_min))
         step_max = max(math.ceil(tol), 1)
-        step_strats.append(xps.from_dtype(dtype, min_value=step_max))
+        step_strats.append(hh.from_dtype(dtype, min_value=step_max))
         step = data.draw(st.one_of(step_strats), label="step")
     assert step != 0, "step must not equal 0"  # sanity check
 
@@ -215,11 +215,11 @@ def test_asarray_scalars(shape, data):
     else:
         _dtype = dtype
     if dh.is_float_dtype(_dtype):
-        elements_strat = xps.from_dtype(_dtype) | xps.from_dtype(xp.int32)
+        elements_strat = hh.from_dtype(_dtype) | hh.from_dtype(xp.int32)
     elif dh.is_int_dtype(_dtype):
-        elements_strat = xps.from_dtype(_dtype) | st.booleans()
+        elements_strat = hh.from_dtype(_dtype) | st.booleans()
     else:
-        elements_strat = xps.from_dtype(_dtype)
+        elements_strat = hh.from_dtype(_dtype)
     size = math.prod(shape)
     obj_strat = st.lists(elements_strat, min_size=size, max_size=size)
     scalar_type = dh.get_scalar_type(_dtype)
@@ -267,7 +267,7 @@ def scalar_eq(s1: Scalar, s2: Scalar) -> bool:
     data=st.data(),
 )
 def test_asarray_arrays(shape, dtypes, data):
-    x = data.draw(xps.arrays(dtype=dtypes.input_dtype, shape=shape), label="x")
+    x = data.draw(hh.arrays(dtype=dtypes.input_dtype, shape=shape), label="x")
     dtypes_strat = st.just(dtypes.input_dtype)
     if dtypes.input_dtype == dtypes.result_dtype:
         dtypes_strat |= st.none()
@@ -290,7 +290,7 @@ def test_asarray_arrays(shape, dtypes, data):
         stype = dh.get_scalar_type(x.dtype)
         idx = data.draw(xps.indices(x.shape, max_dims=0), label="mutating idx")
         old_value = stype(x[idx])
-        scalar_strat = xps.from_dtype(dtypes.input_dtype).filter(
+        scalar_strat = hh.from_dtype(dtypes.input_dtype).filter(
             lambda n: not scalar_eq(n, old_value)
         )
         value = data.draw(
@@ -326,7 +326,7 @@ def test_empty(shape, kw):
 
 
 @given(
-    x=xps.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()),
+    x=hh.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()),
     kw=hh.kwargs(dtype=st.none() | xps.scalar_dtypes()),
 )
 def test_empty_like(x, kw):
@@ -382,7 +382,7 @@ def full_fill_values(draw) -> Union[bool, int, float, complex]:
         st.shared(hh.kwargs(dtype=st.none() | xps.scalar_dtypes()), key="full_kw")
     )
     dtype = kw.get("dtype", None) or draw(default_safe_dtypes)
-    return draw(xps.from_dtype(dtype))
+    return draw(hh.from_dtype(dtype))
 
 
 @given(
@@ -430,8 +430,8 @@ def test_full(shape, fill_value, kw):
 @given(kw=hh.kwargs(dtype=st.none() | xps.scalar_dtypes()), data=st.data())
 def test_full_like(kw, data):
     dtype = kw.get("dtype", None) or data.draw(xps.scalar_dtypes(), label="dtype")
-    x = data.draw(xps.arrays(dtype=dtype, shape=hh.shapes()), label="x")
-    fill_value = data.draw(xps.from_dtype(dtype), label="fill_value")
+    x = data.draw(hh.arrays(dtype=dtype, shape=hh.shapes()), label="x")
+    fill_value = data.draw(hh.from_dtype(dtype), label="fill_value")
     out = xp.full_like(x, fill_value, **kw)
     dtype = kw.get("dtype", None) or x.dtype
     if kw.get("dtype", None) is None:
@@ -454,8 +454,8 @@ def test_full_like(kw, data):
 def test_linspace(num, dtype, endpoint, data):
     _dtype = dh.default_float if dtype is None else dtype
 
-    start = data.draw(xps.from_dtype(_dtype, **finite_kw), label="start")
-    stop = data.draw(xps.from_dtype(_dtype, **finite_kw), label="stop")
+    start = data.draw(hh.from_dtype(_dtype, **finite_kw), label="start")
+    stop = data.draw(hh.from_dtype(_dtype, **finite_kw), label="stop")
     # avoid overflow errors
     assume(not xp.isnan(xp.asarray(stop - start, dtype=_dtype)))
     assume(not xp.isnan(xp.asarray(start - stop, dtype=_dtype)))
@@ -509,7 +509,7 @@ def test_meshgrid(dtype, data):
     )
     arrays = []
     for i, shape in enumerate(shapes, 1):
-        x = data.draw(xps.arrays(dtype=dtype, shape=shape), label=f"x{i}")
+        x = data.draw(hh.arrays(dtype=dtype, shape=shape), label=f"x{i}")
         arrays.append(x)
     # sanity check
     assert math.prod(math.prod(x.shape) for x in arrays) <= hh.MAX_ARRAY_SIZE
@@ -541,7 +541,7 @@ def test_ones(shape, kw):
 
 
 @given(
-    x=xps.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()),
+    x=hh.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()),
     kw=hh.kwargs(dtype=st.none() | xps.scalar_dtypes()),
 )
 def test_ones_like(x, kw):
@@ -579,7 +579,7 @@ def test_zeros(shape, kw):
 
 
 @given(
-    x=xps.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()),
+    x=hh.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()),
     kw=hh.kwargs(dtype=st.none() | xps.scalar_dtypes()),
 )
 def test_zeros_like(x, kw):

array_api_tests/test_data_type_functions.py

@@ -34,7 +34,7 @@ def float32(n: Union[int, float]) -> float:
 )
 def test_astype(x_dtype, dtype, kw, data):
     if xp.bool in (x_dtype, dtype):
-        elements_strat = xps.from_dtype(x_dtype)
+        elements_strat = hh.from_dtype(x_dtype)
     else:
         m1, M1 = dh.dtype_ranges[x_dtype]
         m2, M2 = dh.dtype_ranges[dtype]
@@ -46,15 +46,15 @@ def test_astype(x_dtype, dtype, kw, data):
             cast = float
         min_value = cast(max(m1, m2))
         max_value = cast(min(M1, M2))
-        elements_strat = xps.from_dtype(
+        elements_strat = hh.from_dtype(
             x_dtype,
             min_value=min_value,
             max_value=max_value,
             allow_nan=False,
             allow_infinity=False,
         )
     x = data.draw(
-        xps.arrays(dtype=x_dtype, shape=hh.shapes(), elements=elements_strat), label="x"
+        hh.arrays(dtype=x_dtype, shape=hh.shapes(), elements=elements_strat), label="x"
     )
 
     out = xp.astype(x, dtype, **kw)
@@ -71,7 +71,7 @@ def test_astype(x_dtype, dtype, kw, data):
 def test_broadcast_arrays(shapes, data):
     arrays = []
     for c, shape in enumerate(shapes, 1):
-        x = data.draw(xps.arrays(dtype=xps.scalar_dtypes(), shape=shape), label=f"x{c}")
+        x = data.draw(hh.arrays(dtype=xps.scalar_dtypes(), shape=shape), label=f"x{c}")
         arrays.append(x)
 
     out = xp.broadcast_arrays(*arrays)
@@ -94,7 +94,7 @@ def test_broadcast_arrays(shapes, data):
     # TODO: test values
 
 
-@given(x=xps.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()), data=st.data())
+@given(x=hh.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()), data=st.data())
 def test_broadcast_to(x, data):
     shape = data.draw(
         hh.mutually_broadcastable_shapes(1, base_shape=x.shape)
@@ -113,7 +113,7 @@ def test_broadcast_to(x, data):
 @given(_from=non_complex_dtypes(), to=non_complex_dtypes(), data=st.data())
 def test_can_cast(_from, to, data):
     from_ = data.draw(
-        st.just(_from) | xps.arrays(dtype=_from, shape=hh.shapes()), label="from_"
+        st.just(_from) | hh.arrays(dtype=_from, shape=hh.shapes()), label="from_"
     )
 
     out = xp.can_cast(from_, to)