diff --git a/array_api_tests/dtype_helpers.py b/array_api_tests/dtype_helpers.py index 7e66ca65..ffac46d8 100644 --- a/array_api_tests/dtype_helpers.py +++ b/array_api_tests/dtype_helpers.py @@ -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.""" diff --git a/array_api_tests/hypothesis_helpers.py b/array_api_tests/hypothesis_helpers.py index e6aa1da5..e12d2731 100644 --- a/array_api_tests/hypothesis_helpers.py +++ b/array_api_tests/hypothesis_helpers.py @@ -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 @@ -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) diff --git a/array_api_tests/test_dlpack.py b/array_api_tests/test_dlpack.py index 8251c35d..b2aa4221 100644 --- a/array_api_tests/test_dlpack.py +++ b/array_api_tests/test_dlpack.py @@ -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())