Add JAX support

Unlike other libraries, there is no wrapping for JAX. Actual JAX array_api
support is in JAX itself in the jax.experimental.array_api submodule. This
just adds JAX support to the various helper functions. This also means that we
do not run array-api-tests on JAX.

Closes #83.

Author: asmeurer

README.md

@@ -2,7 +2,7 @@
 
 This is a small wrapper around common array libraries that is compatible with
 the [Array API standard](https://data-apis.org/array-api/latest/). Currently,
-NumPy, CuPy, PyTorch, and Dask are supported. If you want support for other array
+NumPy, CuPy, PyTorch, Dask, and JAX are supported. If you want support for other array
 libraries, or if you encounter any issues, please [open an
 issue](https://github.com/data-apis/array-api-compat/issues).
 
@@ -60,6 +60,12 @@ import array_api_compat.torch as torch
 import array_api_compat.dask as da
 ```
 
+> [!NOTE]
+> There is no `array_api_compat.jax` submodule. JAX support is contained
+> in JAX itself in the `jax.experimental.array_api` module. array-api-compat simply
+> wraps that submodule. The main JAX support in this module consists of
+> supporting it in the [helper functions](#helper-functions) defined below.
+
 Each will include all the functions from the normal NumPy/CuPy/PyTorch/dask.array
 namespace, except that functions that are part of the array API are wrapped so
 that they have the correct array API behavior. In each case, the array object
@@ -104,9 +110,9 @@ part of the specification but which are useful for using the array API:
   object.
 
 - `is_numpy_array(x)`, `is_cupy_array(x)`, `is_torch_array(x)`,
-  `is_dask_array(x)`: return `True` if `x` is an array from the corresponding
-  library. These functions do not import the underlying library if it has not
-  already been imported, so they are cheap to use.
+  `is_dask_array(x)`, `is_jax_array(x)`: return `True` if `x` is an array from
+  the corresponding library. These functions do not import the underlying
+  library if it has not already been imported, so they are cheap to use.
 
 - `array_namespace(*xs)`: Get the corresponding array API namespace for the
   arrays `xs`. For example, if the arrays are NumPy arrays, the returned
@@ -228,6 +234,12 @@ version.
 
 The minimum supported PyTorch version is 1.13.
 
+### JAX
+
+Unlike the other libraries supported here, JAX array API support is contained
+entirely in the JAX library. The JAX array API support is tracked at
+https://github.com/google/jax/issues/18353.
+
 ## Vendoring
 
 This library supports vendoring as an installation method. To vendor the

array_api_compat/common/_helpers.py

@@ -9,6 +9,7 @@
 
 import sys
 import math
+import inspect
 
 def is_numpy_array(x):
     # Avoid importing NumPy if it isn't already
@@ -49,6 +50,15 @@ def is_dask_array(x):
 
     return isinstance(x, dask.array.Array)
 
+def is_jax_array(x):
+    # Avoid importing jax if it isn't already
+    if 'jax' not in sys.modules:
+        return False
+
+    import jax.numpy
+
+    return isinstance(x, jax.numpy.ndarray)
+
 def is_array_api_obj(x):
     """
     Check if x is an array API compatible array object.
@@ -57,6 +67,7 @@ def is_array_api_obj(x):
         or is_cupy_array(x) \
         or is_torch_array(x) \
         or is_dask_array(x) \
+        or is_jax_array(x) \
         or hasattr(x, '__array_namespace__')
 
 def _check_api_version(api_version):
@@ -112,6 +123,13 @@ def your_function(x, y):
                 namespaces.add(dask_namespace)
             else:
                 raise TypeError("_use_compat cannot be False if input array is a dask array!")
+        elif is_jax_array(x):
+            _check_api_version(api_version)
+            # jax.numpy is already an array namespace, but requires this
+            # side-effecting import for __array_namespace__ and some other
+            # things to be defined.
+            import jax.experimental.array_api as jnp
+            namespaces.add(jnp)
         elif hasattr(x, '__array_namespace__'):
             namespaces.add(x.__array_namespace__(api_version=api_version))
         else:
@@ -158,6 +176,15 @@ def device(x: "Array", /) -> "Device":
     """
     if is_numpy_array(x):
         return "cpu"
+    if is_jax_array(x):
+        # JAX has .device() as a method, but it is being deprecated so that it
+        # can become a property, in accordance with the standard. In order for
+        # this function to not break when JAX makes the flip, we check for
+        # both here.
+        if inspect.ismethod(x.device):
+            return x.device()
+        else:
+            return x.device
     return x.device
 
 # Based on cupy.array_api.Array.to_device
@@ -204,6 +231,12 @@ 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``.
@@ -243,6 +276,8 @@ def to_device(x: "Array", device: "Device", /, *, stream: "Optional[Union[int, A
         if device == 'cpu':
             return x
         raise ValueError(f"Unsupported device {device!r}")
+    elif is_jax_array(x):
+        return _jax_to_device(x, device, stream=stream)
     return x.to_device(device, stream=stream)
 
 def size(x):
@@ -255,4 +290,4 @@ def size(x):
 
 __all__ = ['is_array_api_obj', 'array_namespace', 'get_namespace', 'device',
            'to_device', 'size', 'is_numpy_array', 'is_cupy_array',
-           'is_torch_array', 'is_dask_array']
+           'is_torch_array', 'is_dask_array', 'is_jax_array']

tests/test_helpers.py

@@ -1,5 +1,5 @@
 from array_api_compat import (is_numpy_array, is_cupy_array, is_torch_array,
-                              is_dask_array, is_array_api_obj)
+                              is_dask_array, is_jax_array, is_array_api_obj)
 
 from ._helpers import import_
 
@@ -10,6 +10,7 @@
     'cupy': 'is_cupy_array',
     'torch': 'is_torch_array',
     'dask.array': 'is_dask_array',
+    'jax.numpy': 'is_jax_array',
 }
 
 @pytest.mark.parametrize('library', is_functions.keys())

tests/test_isdtype.py

@@ -64,9 +64,12 @@ def isdtype_(dtype_, kind):
     assert type(res) is bool
     return res
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array"])
+@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
 def test_isdtype_spec_dtypes(library):
-    xp = import_('array_api_compat.' + library)
+    if library == "jax.numpy":
+        xp = import_('jax.experimental.array_api')
+    else:
+        xp = import_('array_api_compat.' + library)
 
     isdtype = xp.isdtype
 
@@ -98,10 +101,13 @@ def test_isdtype_spec_dtypes(library):
     'bfloat16',
 ]
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array"])
+@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
 @pytest.mark.parametrize("dtype_", additional_dtypes)
 def test_isdtype_additional_dtypes(library, dtype_):
-    xp = import_('array_api_compat.' + library)
+    if library == "jax.numpy":
+        xp = import_('jax.experimental.array_api')
+    else:
+        xp = import_('array_api_compat.' + library)
 
     isdtype = xp.isdtype