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

@@ -1055,13 +1055,13 @@ def test_array_on_device_with_scalar(self, xp: ModuleType, device: Device):
         b = 1
         res = isclose(a, b)
         assert get_device(res) == device
-        assert_equal(res, xp.asarray([False, False, False, False, True]))
+        assert_equal(res, xp.asarray([False, False, False, False, True], device=device))
 
         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]))
+        assert_equal(res, xp.asarray([False, False, False, False, True], device=device))
 
 
 class TestKron:
@@ -1676,7 +1676,7 @@ def test_assume_unique_and_invert(
 
         a = xp.asarray([0, 3, 6, 10], device=device)
         b = xp.asarray([1, 2, 3, 10], device=device)
-        expected = xp.asarray([True, False, True, False])
+        expected = xp.asarray([True, False, True, False], device=device)
         res = isin(a, b, assume_unique=True, invert=True)
         assert get_device(res) == device
         assert_equal(res, expected)