Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions array_api_tests/dtype_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,13 @@ def get_scalar_type(dtype: DataType) -> ScalarType:
def is_scalar(x):
return isinstance(x, (int, float, complex, bool))

def is_dtype_device_compatible(dtype, device):
try:
xp.asarray([0], dtype=dtype, device=device)
except Exception:
return False
else:
return True

def complex_dtype_for(dtyp):
"""Complex dtype for a float or complex."""
Expand Down
2 changes: 2 additions & 0 deletions array_api_tests/hypothesis_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ def arrays(dtype, *args, elements=None, **kwargs) -> SearchStrategy[Array]:

_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]
_device_dtype_pairs = [(d, device) for d in _sorted_dtypes for device in xp.__array_namespace_info__().devices() if dh.is_dtype_device_compatible(d, device)]

def _dtypes_sorter(dtype_pair: Tuple[DataType, DataType]):
dtype1, dtype2 = dtype_pair
Expand Down Expand Up @@ -199,6 +200,7 @@ def oneway_broadcastable_shapes(draw) -> OnewayBroadcastableShapes:
# Use these instead of xps.scalar_dtypes, etc. because it skips dtypes from
# ARRAY_API_TESTS_SKIP_DTYPES
all_dtypes = sampled_from(_sorted_dtypes)
device_dtype_pairs = sampled_from(_device_dtype_pairs)
all_int_dtypes = sampled_from(dh.all_int_dtypes)
int_dtypes = sampled_from(dh.int_dtypes) # signed ints
uint_dtypes = sampled_from(dh.uint_dtypes)
Expand Down
10 changes: 5 additions & 5 deletions array_api_tests/test_dlpack.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,22 +88,22 @@ def test_dunder_dlpack(x, copy_kw, max_version_kw, dl_device_kw, data):


@given(
x=hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_dims=1, max_side=2)),
dtype_device_pair = hh.device_dtype_pairs,
copy_kw=hh.kwargs(copy=st.booleans()),
data=st.data()
)
def test_from_dlpack(x, copy_kw, data):
def test_from_dlpack(copy_kw, data, dtype_device_pair):
# TODO: 1. test copy; 2. generate inputs on non-default devices;
# 3. test for copy=False cross-device transfers
# 4. test 0D arrays / numpy scalars (the latter do not support dlpack ATM)

dtype, device = dtype_device_pair
x = data.draw(hh.arrays(dtype=dtype, shape=hh.shapes(min_dims=1, max_side=2)))
copy = copy_kw["copy"] if copy_kw else None
if copy is False:
# XXX there is no way to tell if a no-copy cross-device transfer is meant to succeed
devices = [x.device]
else:
devices = xp.__array_namespace_info__().devices()
devices = _compatible_devices(devices)
devices = [device]

tgt_device_kw = data.draw(
hh.kwargs(device=st.sampled_from(devices) | st.none())
Expand Down
Loading