Add basic JAX support

array_api_compat/common/_helpers.py

@@ -55,9 +55,9 @@ def is_jax_array(x):
     if 'jax' not in sys.modules:
         return False
 
-    import jax.numpy
+    import jax
 
-    return isinstance(x, jax.numpy.ndarray)
+    return isinstance(x, jax.Array)
 
 def is_array_api_obj(x):
     """
@@ -153,7 +153,7 @@ def _check_device(xp, device):
         if device not in ["cpu", None]:
             raise ValueError(f"Unsupported device for NumPy: {device!r}")
 
-# device() is not on numpy.ndarray and and to_device() is not on numpy.ndarray
+# device() is not on numpy.ndarray and to_device() is not on numpy.ndarray
 # or cupy.ndarray. They are not included in array objects of this library
 # because this library just reuses the respective ndarray classes without
 # wrapping or subclassing them. These helper functions can be used instead of
@@ -230,12 +230,6 @@ def _torch_to_device(x, device, /, stream=None):
         raise NotImplementedError
     return x.to(device)
 
-def _jax_to_device(x, device, /, stream=None):
-    import jax
-    if stream is not None:
-        raise NotImplementedError
-    return jax.device_put(x, device)
-
 def to_device(x: "Array", device: "Device", /, *, stream: "Optional[Union[int, Any]]" = None) -> "Array":
     """
     Copy the array from the device on which it currently resides to the specified ``device``.
@@ -276,7 +270,9 @@ def to_device(x: "Array", device: "Device", /, *, stream: "Optional[Union[int, A
             return x
         raise ValueError(f"Unsupported device {device!r}")
     elif is_jax_array(x):
-        return _jax_to_device(x, device, stream=stream)
+        # This import adds to_device to x
+        import jax.experimental.array_api
+        return x.to_device(device, stream=stream)
     return x.to_device(device, stream=stream)
 
 def size(x):

tests/test_helpers.py

@@ -1,5 +1,6 @@
 from array_api_compat import (is_numpy_array, is_cupy_array, is_torch_array,
-                              is_dask_array, is_jax_array, is_array_api_obj)
+                              is_dask_array, is_jax_array, is_array_api_obj,
+                              device, to_device)
 
 from ._helpers import import_
 
@@ -24,3 +25,19 @@ def test_is_xp_array(library, func):
     assert is_func(x) == (func == is_functions[library])
 
     assert is_array_api_obj(x)
+
+@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+def test_device(library):
+    if library == "jax.numpy":
+        xp = import_('jax.experimental.array_api')
+    else:
+        xp = import_('array_api_compat.' + library)
+
+    # We can't test much for device() and to_device() other than that
+    # x.to_device(x.device) works.
+
+    x = xp.asarray([1, 2, 3])
+    dev = device(x)
+
+    x2 = to_device(x, dev)
+    assert device(x) == device(x2)