test the right way

Author: lithomas1

src/array_api_extra/_lib/_testing.py

@@ -10,7 +10,6 @@
 from typing import cast
 
 import pytest
-from array_api_compat import is_array_api_strict_namespace
 
 from ._utils._compat import (
     array_namespace,
@@ -107,11 +106,6 @@ def xp_assert_equal(actual: Array, desired: Array, err_msg: str = "") -> None:
             desired = desired.todense()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
 
         # JAX uses `np.testing`
-        if is_array_api_strict_namespace(xp):
-            # Have to move to CPU for array API strict devices before
-            # we're allowed to convert into numpy
-            actual = np.asarray(xp.asarray(actual, device=xp.Device("CPU_DEVICE")))
-            desired = np.asarray(xp.asarray(actual, device=xp.Device("CPU_DEVICE")))
         np.testing.assert_array_equal(actual, desired, err_msg=err_msg)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 
 

tests/conftest.py

@@ -113,7 +113,7 @@ def xp(
     The current array namespace.
     """
     if library == Backend.NUMPY_READONLY:
-        return NumPyReadOnly(), None  # type: ignore[return-value]  # pyright: ignore[reportReturnType]
+        return NumPyReadOnly()  # type: ignore[return-value]  # pyright: ignore[reportReturnType]
     xp = pytest.importorskip(library.value)
     # Possibly wrap module with array_api_compat
     xp = array_namespace(xp.empty(0))
@@ -131,11 +131,7 @@ def xp(
         # suppress unused-ignore to run mypy in -e lint as well as -e dev
         jax.config.update("jax_enable_x64", True)  # type: ignore[no-untyped-call,unused-ignore]
 
-    device = None
-    if library == Backend.ARRAY_API_STRICT_DEVICE1:
-        import array_api_strict
-        device = array_api_strict.Device("device1")
-    return xp, device
+    return xp
 
 
 @pytest.fixture(params=[Backend.DASK])  # Can select the test with `pytest -k dask`

tests/test_funcs.py

@@ -633,9 +633,8 @@ def test_some_inf(self, xp: ModuleType):
         xp_assert_equal(actual, xp.asarray([True, True, True, False, False]))
 
     def test_equal_nan(self, xp: ModuleType):
-        xp, device = xp
-        a = xp.asarray([xp.nan, xp.nan, 1.0], device=device)
-        b = xp.asarray([xp.nan, 1.0, xp.nan], device=device)
+        a = xp.asarray([xp.nan, xp.nan, 1.0])
+        b = xp.asarray([xp.nan, 1.0, xp.nan])
         xp_assert_equal(isclose(a, b), xp.asarray([False, False, False]))
         xp_assert_equal(isclose(a, b, equal_nan=True), xp.asarray([True, False, False]))
 
@@ -721,6 +720,11 @@ def test_xp(self, xp: ModuleType):
         b = xp.asarray([1e-9, 1e-4])
         xp_assert_equal(isclose(a, b, xp=xp), xp.asarray([True, False]))
 
+    @pytest.mark.parametrize("equal_nan", [True, False])
+    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)
+        xp_assert_equal(isclose(a, b, equal_nan=equal_nan), xp.asarray([True, False, equal_nan]))
 
 class TestKron:
     def test_basic(self, xp: ModuleType):