Test shape and dtype for xp.asarray()

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -18,7 +18,7 @@
 from .algos import broadcast_shapes
 from .function_stubs import elementwise_functions
 from .pytest_helpers import nargs
-from .typing import Array, DataType, Scalar, Shape
+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
@@ -430,11 +430,3 @@ def axes(ndim: int) -> SearchStrategy[Optional[Union[int, Shape]]]:
         axes_strats.append(integers(-ndim, ndim - 1))
         axes_strats.append(xps.valid_tuple_axes(ndim))
     return one_of(axes_strats)
-
-
-def scalar_objects(dtype: DataType, shape: Shape) -> SearchStrategy[List[Scalar]]:
-    """Generates scalars or nested sequences which are valid for xp.asarray()"""
-    size = math.prod(shape)
-    return lists(xps.from_dtype(dtype), min_size=size, max_size=size).map(
-        lambda l: sh.reshape(l, shape)
-    )

array_api_tests/test_array_object.py

@@ -1,6 +1,6 @@
 import math
 from itertools import product
-from typing import get_args
+from typing import List, get_args
 
 import pytest
 from hypothesis import assume, given, note
@@ -12,13 +12,21 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from .typing import DataType, Param, Scalar, ScalarType
+from .typing import DataType, Param, Scalar, ScalarType, Shape
+
+
+def scalar_objects(dtype: DataType, shape: Shape) -> st.SearchStrategy[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(
+        lambda l: sh.reshape(l, shape)
+    )
 
 
 @given(hh.shapes(min_side=1), st.data())  # TODO: test 0-sided arrays
 def test_getitem(shape, data):
     dtype = data.draw(xps.scalar_dtypes(), label="dtype")
-    obj = data.draw(hh.scalar_objects(dtype, shape), label="obj")
+    obj = data.draw(scalar_objects(dtype, shape), label="obj")
     x = xp.asarray(obj, dtype=dtype)
     note(f"{x=}")
     key = data.draw(xps.indices(shape=shape), label="key")
@@ -59,7 +67,7 @@ def test_getitem(shape, data):
 @given(hh.shapes(min_side=1), st.data())  # TODO: test 0-sided arrays
 def test_setitem(shape, data):
     dtype = data.draw(xps.scalar_dtypes(), label="dtype")
-    obj = data.draw(hh.scalar_objects(dtype, shape), label="obj")
+    obj = data.draw(scalar_objects(dtype, shape), label="obj")
     x = xp.asarray(obj, dtype=dtype)
     note(f"{x=}")
     # TODO: test setting non-0d arrays

array_api_tests/test_creation_functions.py

@@ -10,6 +10,7 @@
 from . import dtype_helpers as dh
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
+from . import shape_helpers as sh
 from . import xps
 from .typing import DataType, Scalar
 
@@ -186,14 +187,59 @@ def test_arange(dtype, data):
             ), f"out[0]={out[0]}, but should be {_start} {f_func}"
 
 
-@given(dtype=xps.scalar_dtypes(), shape=hh.shapes(min_side=1), data=st.data())
-def test_asarray_scalars(dtype, shape, data):
-    obj = data.draw(hh.scalar_objects(dtype, shape), label="obj")
+@given(
+    shape=hh.shapes(min_side=1),
+    data=st.data(),
+)
+def test_asarray_scalars(shape, data):
     kw = data.draw(
-        hh.kwargs(dtype=st.sampled_from([None, dtype]), copy=st.none()), label="kw"
+        hh.kwargs(dtype=st.none() | xps.scalar_dtypes(), copy=st.none()), label="kw"
     )
+    dtype = kw.get("dtype", None)
+    if dtype is None:
+        dtype_family = data.draw(
+            st.sampled_from(
+                [(xp.bool,), (xp.int32, xp.int64), (xp.float32, xp.float64)]
+            ),
+            label="expected out dtypes",
+        )
+        _dtype = dtype_family[0]
+    else:
+        _dtype = dtype
+    if dh.is_float_dtype(_dtype):
+        elements_strat = xps.from_dtype(_dtype) | xps.from_dtype(xp.int32)
+    elif dh.is_int_dtype(_dtype):
+        elements_strat = xps.from_dtype(_dtype) | st.booleans()
+    else:
+        elements_strat = xps.from_dtype(_dtype)
+    size = math.prod(shape)
+    obj_strat = st.lists(elements_strat, min_size=size, max_size=size)
+    if dtype is None:
+        # For asarray to infer the dtype we're testing, obj requires at least
+        # one element to be the scalar equivalent of the inferred dtype, and so
+        # we filter out invalid examples. Note we use type() as Python booleans
+        # instance check with ints e.g. isinstance(False, int) == True.
+        scalar_type = dh.get_scalar_type(_dtype)
+        obj_strat = obj_strat.filter(lambda l: any(type(e) == scalar_type for e in l))
+    obj_strat = obj_strat.map(lambda l: sh.reshape(l, shape))
+    obj = data.draw(obj_strat, label="obj")
+
+    out = xp.asarray(obj, **kw)
 
-    xp.asarray(obj, **kw)
+    if dtype is None:
+        msg = f"out.dtype={dh.dtype_to_name[out.dtype]}, should be "
+        if dtype_family == (xp.float32, xp.float64):
+            msg += "default floating-point dtype (float32 or float64)"
+        elif dtype_family == (xp.int32, xp.int64):
+            msg += "default integer dtype (int32 or int64)"
+        else:
+            msg += "boolean dtype"
+        msg += " [asarray()]"
+        assert out.dtype in dtype_family, msg
+    else:
+        assert kw["dtype"] == _dtype  # sanity check
+        ph.assert_kw_dtype("asarray", _dtype, out.dtype)
+    ph.assert_shape("asarray", out.shape, shape)
 
 
 # TODO: test asarray with arrays and copy (in a seperate method)