convert to CPU in xp_assert_equal

Author: lithomas1

src/array_api_extra/_lib/_testing.py

@@ -13,6 +13,7 @@
 
 from ._utils._compat import (
     array_namespace,
+    is_array_api_strict_namespace,
     is_cupy_namespace,
     is_dask_namespace,
     is_pydata_sparse_namespace,
@@ -105,6 +106,12 @@ def xp_assert_equal(actual: Array, desired: Array, err_msg: str = "") -> None:
             actual = actual.todense()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
             desired = desired.todense()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
 
+        if is_array_api_strict_namespace(xp):
+            # __array__ doesn't work on array-api-strict device arrays
+            # We need to convert to the CPU device first
+            actual = np.asarray(xp.asarray(actual, device=xp.Device("CPU_DEVICE")))
+            desired = np.asarray(xp.asarray(actual, device=xp.Device("CPU_DEVICE")))
+
         # JAX uses `np.testing`
         np.testing.assert_array_equal(actual, desired, err_msg=err_msg)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 

tests/test_funcs.py

@@ -724,6 +724,8 @@ def test_xp(self, xp: ModuleType):
     def test_device(self, xp: ModuleType, device: Device, equal_nan: bool):
         a = xp.asarray([0.0, 0.0, xp.nan], device=device)
         b = xp.asarray([1e-9, 1e-4, xp.nan], device=device)
+        res = isclose(a, b, equal_nan=equal_nan)
+        assert get_device(res) == device
         xp_assert_equal(
             isclose(a, b, equal_nan=equal_nan), xp.asarray([True, False, equal_nan])
         )