Hack around JAX device attribute

Author: crusaderky

src/array_api_extra/_lib/_utils/_compat.py

@@ -1,11 +1,17 @@
 """Acquire helpers from array-api-compat."""
+
 # Allow packages that vendor both `array-api-extra` and
 # `array-api-compat` to override the import location
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ._typing import Array, Device
 
 try:
     from ...._array_api_compat_vendor import (
         array_namespace,
-        device,
         is_array_api_strict_namespace,
         is_cupy_namespace,
         is_dask_namespace,
@@ -17,10 +23,12 @@
         is_writeable_array,
         size,
     )
+    from ...._array_api_compat_vendor import (
+        device as _compat_device,
+    )
 except ImportError:
     from array_api_compat import (
         array_namespace,
-        device,
         is_array_api_strict_namespace,
         is_cupy_namespace,
         is_dask_namespace,
@@ -32,6 +40,14 @@
         is_writeable_array,
         size,
     )
+    from array_api_compat import (
+        device as _compat_device,
+    )
+
+
+def device(x: Array) -> Device | None:  # numpydoc ignore=GL08
+    return None if is_jax_array(x) else _compat_device(x)
+
 
 __all__ = [
     "array_namespace",