ENH: testing: Add check_device parameter

src/array_api_extra/testing.py

@@ -28,6 +28,9 @@
     is_torch_namespace,
     to_device,
 )
+from ._lib._utils._compat import (
+    device as get_device,
+)
 from ._lib._utils._helpers import jax_autojit, pickle_flatten, pickle_unflatten
 from ._lib._utils._typing import Array, Device
 
@@ -564,12 +567,13 @@ def _require_numpy() -> ModuleType:  # numpydoc ignore=RT01
     return np
 
 
-def _check_ns_shape_dtype(
+def _check_ns_shape_dtype_device(
     actual: Array,
     desired: Array,
     check_dtype: bool,
     check_shape: bool,
     check_scalar: bool,
+    check_device: bool,
     xp: ModuleType | None = None,
 ) -> tuple[Array, Array, ModuleType, ModuleType]:  # numpydoc ignore=RT03
     """
@@ -588,6 +592,8 @@ def _check_ns_shape_dtype(
     check_scalar : bool, default: False
         NumPy only: whether to check agreement between actual and desired types -
         0d array vs scalar.
+    check_device : bool, default: True
+        Whether to check agreement between actual and desired devices.
     xp : array_namespace, optional
         A standard-compatible namespace which `actual` and `desired` must match.
 
@@ -655,6 +661,12 @@ def _check_ns_shape_dtype(
     if check_dtype:
         msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}"
         assert actual.dtype == desired.dtype, msg
+    if check_device:
+        msg = (
+            f"Devices do not match.\nActual: {get_device(actual)}\n"
+            f"Desired: {get_device(desired)}"
+        )
+        assert get_device(actual) == get_device(desired), msg
     desired = desired_xp.broadcast_to(desired, actual_shape)
     return actual, desired, desired_xp, np
 
@@ -714,6 +726,7 @@ def assert_close(
     check_dtype: bool = True,
     check_shape: bool = True,
     check_scalar: bool = False,
+    check_device: bool = True,
     xp: ModuleType | None = None,
 ) -> None:
     """
@@ -746,6 +759,8 @@ def assert_close(
     check_scalar : bool, default: False
         NumPy only: whether to check agreement between actual and desired types —
         0-D :class:`numpy.ndarray` vs scalar (e.g. :class:`numpy.double`).
+    check_device : bool, default: True
+        Whether to check agreement between actual and desired devices.
     xp : array_namespace, optional
         A standard-compatible namespace which `actual` and `desired` must match.
 
@@ -777,8 +792,8 @@ def assert_close(
     Array arguments to `atol` and `rtol` must be valid input to :class:`float`.
     """
     __tracebackhide__ = True
-    actual, desired, xp, np = _check_ns_shape_dtype(
-        actual, desired, check_dtype, check_shape, check_scalar, xp
+    actual, desired, xp, np = _check_ns_shape_dtype_device(
+        actual, desired, check_dtype, check_shape, check_scalar, check_device, xp
     )
     if not _is_materializable(actual):
         return
@@ -818,6 +833,7 @@ def assert_equal(
     check_dtype: bool = True,
     check_shape: bool = True,
     check_scalar: bool = False,
+    check_device: bool = True,
     xp: ModuleType | None = None,
 ) -> None:
     """
@@ -844,6 +860,8 @@ def assert_equal(
     check_scalar : bool, default: False
         NumPy only: whether to check agreement between actual and desired types —
         0-D :class:`numpy.ndarray` vs scalar (e.g. :class:`numpy.double`).
+    check_device : bool, default: True
+        Whether to check agreement between actual and desired devices.
     xp : array_namespace, optional
         A standard-compatible namespace which `actual` and `desired` must match.
 
@@ -861,8 +879,8 @@ def assert_equal(
     numpy.testing.assert_array_equal : Similar function for NumPy arrays.
     """
     __tracebackhide__ = True
-    actual, desired, xp, np = _check_ns_shape_dtype(
-        actual, desired, check_dtype, check_shape, check_scalar, xp
+    actual, desired, xp, np = _check_ns_shape_dtype_device(
+        actual, desired, check_dtype, check_shape, check_scalar, check_device, xp
     )
     if not _is_materializable(actual):
         return
@@ -882,6 +900,7 @@ def assert_less(
     check_dtype: bool = True,
     check_shape: bool = True,
     check_scalar: bool = False,
+    check_device: bool = True,
     xp: ModuleType | None = None,
 ) -> None:
     """
@@ -906,6 +925,8 @@ def assert_less(
     check_scalar : bool, default: False
         NumPy only: whether to check agreement between actual and desired types —
         0-D :class:`numpy.ndarray` vs scalar (e.g. :class:`numpy.double`).
+    check_device : bool, default: True
+        Whether to check agreement between actual and desired devices.
     xp : array_namespace, optional
         A standard-compatible namespace which `x` and `y` must match.
 
@@ -923,8 +944,8 @@ def assert_less(
     numpy.testing.assert_array_less : Similar function for NumPy arrays.
     """
     __tracebackhide__ = True
-    x, y, xp, np = _check_ns_shape_dtype(
-        x, y, check_dtype, check_shape, check_scalar, xp
+    x, y, xp, np = _check_ns_shape_dtype_device(
+        x, y, check_dtype, check_shape, check_scalar, check_device, xp
     )
     if not _is_materializable(x):
         return
@@ -941,6 +962,7 @@ def assert_close_nulp(
     check_dtype: bool = True,
     check_shape: bool = True,
     check_scalar: bool = False,
+    check_device: bool = True,
     xp: ModuleType | None = None,
 ) -> None:
     """
@@ -966,6 +988,8 @@ def assert_close_nulp(
     check_scalar : bool, default: False
         NumPy only: whether to check agreement between actual and desired types —
         0-D :class:`numpy.ndarray` vs scalar (e.g. :class:`numpy.double`).
+    check_device : bool, default: True
+        Whether to check agreement between actual and desired devices.
     xp : array_namespace, optional
         A standard-compatible namespace which `actual` and `desired` must match.
 
@@ -996,8 +1020,8 @@ def assert_close_nulp(
     where ``spacing(x)`` is the distance between ``x`` and the nearest adjacent number
     representable by in the data type of ``x``.
     """
-    actual, desired, xp, np = _check_ns_shape_dtype(
-        actual, desired, check_dtype, check_shape, check_scalar, xp
+    actual, desired, xp, np = _check_ns_shape_dtype_device(
+        actual, desired, check_dtype, check_shape, check_scalar, check_device, xp
     )
     if not _is_materializable(actual):
         return

tests/test_funcs.py

@@ -211,15 +211,6 @@ def test_xp(self, xp: ModuleType):
         expect = xp.where(cond, self.f1(x), self.f2(x))
         assert_equal(actual, expect)
 
-    def test_device(self, xp: ModuleType, device: Device):
-        x = xp.asarray([1, 2, 3, 4], device=device)
-        y = apply_where(x % 2 == 0, x, self.f1, self.f2)
-        assert get_device(y) == device
-        y = apply_where(x % 2 == 0, x, self.f1, fill_value=0)
-        assert get_device(y) == device
-        y = apply_where(x % 2 == 0, x, self.f1, fill_value=x)
-        assert get_device(y) == device
-
     @pytest.mark.filterwarnings("ignore::RuntimeWarning")  # overflows, etc.
     @hypothesis.settings(
         # The xp and library fixtures are not regenerated between hypothesis iterations
@@ -1050,19 +1041,6 @@ def test_device(self, xp: ModuleType, device: Device, equal_nan: bool):
         res = isclose(a, b, equal_nan=equal_nan)
         assert get_device(res) == device
 
-    def test_array_on_device_with_scalar(self, xp: ModuleType, device: Device):
-        a = xp.asarray([0.01, 0.5, 0.8, 0.9, 1.00001], device=device, dtype=xp.float64)
-        b = 1
-        res = isclose(a, b)
-        assert get_device(res) == device
-        assert_equal(res, xp.asarray([False, False, False, False, True]))
-
-        a = 0.1
-        b = xp.asarray([0.01, 0.5, 0.8, 0.9, 0.100001], device=device, dtype=xp.float64)
-        res = isclose(a, b)
-        assert get_device(res) == device
-        assert_equal(res, xp.asarray([False, False, False, False, True]))
-
 
 class TestKron:
     def test_basic(self, xp: ModuleType):
@@ -1668,17 +1646,14 @@ def test_device(self, xp: ModuleType, device: Device, library: Backend):
         b = xp.asarray([1, 2, 3], device=device)
         assert get_device(isin(a, b)) == device
 
-    def test_assume_unique_and_invert(
-        self, xp: ModuleType, device: Device, library: Backend
-    ):
+    def test_assume_unique_and_invert(self, xp: ModuleType, library: Backend):
         if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24):
             pytest.xfail("NumPy <1.24 has no kind kwarg in isin")
 
-        a = xp.asarray([0, 3, 6, 10], device=device)
-        b = xp.asarray([1, 2, 3, 10], device=device)
+        a = xp.asarray([0, 3, 6, 10])
+        b = xp.asarray([1, 2, 3, 10])
         expected = xp.asarray([True, False, True, False])
         res = isin(a, b, assume_unique=True, invert=True)
-        assert get_device(res) == device
         assert_equal(res, expected)
 
     def test_kind(self, xp: ModuleType, library: Backend):

tests/test_testing.py

@@ -13,6 +13,9 @@
     is_dask_namespace,
     is_jax_namespace,
 )
+from array_api_extra._lib._utils._compat import (
+    device as get_device,
+)
 from array_api_extra._lib._utils._typing import Array, Device
 from array_api_extra.testing import (
     _as_numpy_array,
@@ -206,13 +209,25 @@ def test_device(self, xp: ModuleType, device: Device, func: Callable[..., None])
         b = xp.asarray([2], device=device)
         c = xp.asarray([2, 2], device=device)
 
-        func(a, b)
+        func(a, b, check_device=True)
         with pytest.raises(AssertionError, match="shapes do not match"):
             func(a, c)
         # This is normally performed by np.testing.assert_array_equal etc.
         # but in case of torch device='meta' we have to do it manually
         with pytest.raises(AssertionError, match="sizes do not match"):
-            func(a, c, check_shape=False)
+            func(a, c, check_shape=False, check_device=False)
+
+        # Checks for non-default device
+        default_device = get_device(xp.empty(0))
+        if device != default_device:
+            d = xp.asarray([2], device=default_device)
+            with pytest.raises(AssertionError, match="Devices do not match"):
+                func(a, d)
+            func(a, d, check_device=False)
+            if getattr(device, "type", None) != "meta":
+                e = xp.asarray([0], device=default_device)
+                with pytest.raises(AssertionError, match=self.np_err_msg(func)):
+                    func(a, e, check_device=False)
 
     def test_assert_close_nulp(self, xp: ModuleType):
         a = xp.asarray([1.0, 1e-10], dtype=xp.float64)