ENH: Generation of compatable device and dtype pairs for test_from_dlpack (#443)

reviewed at #443

Author: prady0t

array_api_tests/dtype_helpers.py

@@ -199,6 +199,31 @@ def is_scalar(x):
     return isinstance(x, (int, float, complex, bool))
 
 
+def is_device_dlpack_compatible(device):
+    """If device is dlpack compatible, return True, else False"""
+    # XXX: there seems to be no better way than try-catch for __dlpack_device__()
+
+    x = xp.empty(2, device=device)
+    try:
+        x.__dlpack_device__()
+    except:
+        # case in point: torch.device(type="meta") raises
+        # ValueError: Unknown device type meta for Dlpack
+        return False
+    else:
+        # no exception => device is compatible (or a cuda device)
+        return True
+
+
+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."""
     if dtyp in complex_dtypes:

array_api_tests/hypothesis_helpers.py

@@ -93,6 +93,13 @@ 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 = [
+    (dtype, device)
+    for dtype in _sorted_dtypes
+    for device in xp.__array_namespace_info__().devices()
+    if dh.is_device_dlpack_compatible(device)
+    if dh.is_dtype_device_compatible(dtype, device)
+]
 
 def _dtypes_sorter(dtype_pair: Tuple[DataType, DataType]):
     dtype1, dtype2 = dtype_pair
@@ -199,6 +206,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)

array_api_tests/test_dlpack.py

@@ -21,29 +21,6 @@ class DLPackDeviceEnum(Enum):
     ONE_API = 14
 
 
-def _compatible_devices(devices):
-    """Given a list of devices, filter out dlpack-incompatible ones."""
-    # XXX: there seems to be no better way than try-catch for __dlpack_device__()
-
-    # XXX: this process actually fails with CuPy because CuPy ignores the device= argument
-    # cf https://github.com/data-apis/array-api-compat/issues/337 and
-    # https://github.com/cupy/cupy/issues/9848
-    # Luckily, CuPy only supports CUDA devices, and they are all compatible.
-    compatible_ = []
-    for device in devices:
-        x = xp.empty(2, device=device)
-        try:
-            x.__dlpack_device__()
-        except:
-            # case in point: torch.device(type="meta") raises
-            # ValueError: Unknown device type meta for Dlpack
-            pass
-        else:
-            # no exception => device is compatible
-            compatible_.append(device)
-    return compatible_
-
-
 @given(dtype=hh.all_dtypes, data=st.data())
 def test_dlpack_device(dtype, data):
     """Test the array object __dlpack_device__ method."""
@@ -88,27 +65,24 @@ 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]
+    # XXX there is no way to tell if a no-copy cross-device transfer is meant to succeed
+    tgt_device = x.device if copy is False else device
+    if  data.draw(st.booleans()):
+        tgt_device_kw = {"device": tgt_device}
     else:
-        devices = xp.__array_namespace_info__().devices()
-        devices = _compatible_devices(devices)
-
-    tgt_device_kw = data.draw(
-        hh.kwargs(device=st.sampled_from(devices) | st.none())
-    )
-    tgt_device = tgt_device_kw['device'] if tgt_device_kw else None
+        tgt_device_kw = {}
+    tgt_device = tgt_device_kw.get("device")
 
     repro_snippet = ph.format_snippet(
         f"y = from_dlpack({x!r}, **tgt_device_kw, **copy_kw) with {tgt_device_kw=} and {copy_kw=}"