BUG: Fix isclose multidevice

Author: lithomas1

src/array_api_extra/_lib/_backends.py

@@ -24,6 +24,7 @@ class Backend(Enum):  # numpydoc ignore=PR01,PR02  # type: ignore[no-subclass-an
     """
 
     ARRAY_API_STRICT = "array_api_strict", _compat.is_array_api_strict_namespace
+    ARRAY_API_STRICT_DEVICE1 = "array_api_strict", _compat.is_array_api_strict_namespace
     NUMPY = "numpy", _compat.is_numpy_namespace
     NUMPY_READONLY = "numpy_readonly", _compat.is_numpy_namespace
     CUPY = "cupy", _compat.is_cupy_namespace

src/array_api_extra/_lib/_funcs.py

@@ -549,7 +549,7 @@ def isclose(
             xp=xp,
         )
         if equal_nan:
-            out = xp.where(xp.isnan(a) & xp.isnan(b), xp.asarray(True), out)
+            out = xp.where(xp.isnan(a) & xp.isnan(b), True, out)
         return out
 
     if xp.isdtype(a.dtype, "bool") or xp.isdtype(b.dtype, "bool"):

src/array_api_extra/_lib/_testing.py

@@ -10,6 +10,7 @@
 from typing import cast
 
 import pytest
+from array_api_compat import is_array_api_strict_namespace
 
 from ._utils._compat import (
     array_namespace,
@@ -106,6 +107,11 @@ 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()  # type: ignore[return-value]  # pyright: ignore[reportReturnType]
+        return NumPyReadOnly(), None  # 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,7 +131,11 @@ 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]
 
-    return xp
+    device = None
+    if library == Backend.ARRAY_API_STRICT_DEVICE1:
+        import array_api_strict
+        device = array_api_strict.Device("device1")
+    return xp, device
 
 
 @pytest.fixture(params=[Backend.DASK])  # Can select the test with `pytest -k dask`

tests/test_funcs.py

@@ -633,8 +633,9 @@ 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):
-        a = xp.asarray([xp.nan, xp.nan, 1.0])
-        b = xp.asarray([xp.nan, 1.0, xp.nan])
+        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)
         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]))