Add basic JAX support

.github/workflows/tests.yml

@@ -15,7 +15,7 @@ jobs:
       - name: Install Dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest numpy torch dask[array]
+          python -m pip install pytest numpy torch dask[array] jax
 
       - name: Run Tests
         run: |

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, and PyTorch 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).
 
@@ -56,7 +56,17 @@ import array_api_compat.cupy as cp
 import array_api_compat.torch as torch
 ```
 
-Each will include all the functions from the normal NumPy/CuPy/PyTorch
+```py
+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
 used will be the same array object from the wrapped library.
@@ -99,6 +109,11 @@ part of the specification but which are useful for using the array API:
 - `is_array_api_obj(x)`: Return `True` if `x` is an array API compatible array
   object.
 
+- `is_numpy_array(x)`, `is_cupy_array(x)`, `is_torch_array(x)`,
+  `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
   namespace will be `array_api_compat.numpy`. Note that this function will
@@ -219,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/_aliases.py

@@ -13,7 +13,7 @@
 from types import ModuleType
 import inspect
 
-from ._helpers import _check_device, _is_numpy_array, array_namespace
+from ._helpers import _check_device, is_numpy_array, array_namespace
 
 # These functions are modified from the NumPy versions.
 
@@ -309,7 +309,7 @@ def _asarray(
         raise ValueError("Unrecognized namespace argument to asarray()")
 
     _check_device(xp, device)
-    if _is_numpy_array(obj):
+    if is_numpy_array(obj):
         import numpy as np
         if hasattr(np, '_CopyMode'):
             # Not present in older NumPys

array_api_compat/common/_helpers.py

@@ -9,8 +9,9 @@
 
 import sys
 import math
+import inspect
 
-def _is_numpy_array(x):
+def is_numpy_array(x):
     # Avoid importing NumPy if it isn't already
     if 'numpy' not in sys.modules:
         return False
@@ -20,7 +21,7 @@ def _is_numpy_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, (np.ndarray, np.generic))
 
-def _is_cupy_array(x):
+def is_cupy_array(x):
     # Avoid importing NumPy if it isn't already
     if 'cupy' not in sys.modules:
         return False
@@ -30,7 +31,7 @@ def _is_cupy_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, (cp.ndarray, cp.generic))
 
-def _is_torch_array(x):
+def is_torch_array(x):
     # Avoid importing torch if it isn't already
     if 'torch' not in sys.modules:
         return False
@@ -40,7 +41,7 @@ def _is_torch_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, torch.Tensor)
 
-def _is_dask_array(x):
+def is_dask_array(x):
     # Avoid importing dask if it isn't already
     if 'dask.array' not in sys.modules:
         return False
@@ -49,14 +50,24 @@ 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.
     """
-    return _is_numpy_array(x) \
-        or _is_cupy_array(x) \
-        or _is_torch_array(x) \
-        or _is_dask_array(x) \
+    return is_numpy_array(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):
@@ -81,37 +92,44 @@ def your_function(x, y):
     """
     namespaces = set()
     for x in xs:
-        if _is_numpy_array(x):
+        if is_numpy_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from .. import numpy as numpy_namespace
                 namespaces.add(numpy_namespace)
             else:
                 import numpy as np
                 namespaces.add(np)
-        elif _is_cupy_array(x):
+        elif is_cupy_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from .. import cupy as cupy_namespace
                 namespaces.add(cupy_namespace)
             else:
                 import cupy as cp
                 namespaces.add(cp)
-        elif _is_torch_array(x):
+        elif is_torch_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from .. import torch as torch_namespace
                 namespaces.add(torch_namespace)
             else:
                 import torch
                 namespaces.add(torch)
-        elif _is_dask_array(x):
+        elif is_dask_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from ..dask import array as dask_namespace
                 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:
@@ -156,8 +174,17 @@ def device(x: "Array", /) -> "Device":
     out: device
         a ``device`` object (see the "Device Support" section of the array API specification).
     """
-    if _is_numpy_array(x):
+    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``.
@@ -225,24 +258,26 @@ def to_device(x: "Array", device: "Device", /, *, stream: "Optional[Union[int, A
     .. note::
        If ``stream`` is given, the copy operation should be enqueued on the provided ``stream``; otherwise, the copy operation should be enqueued on the default stream/queue. Whether the copy is performed synchronously or asynchronously is implementation-dependent. Accordingly, if synchronization is required to guarantee data safety, this must be clearly explained in a conforming library's documentation.
     """
-    if _is_numpy_array(x):
+    if is_numpy_array(x):
         if stream is not None:
             raise ValueError("The stream argument to to_device() is not supported")
         if device == 'cpu':
             return x
         raise ValueError(f"Unsupported device {device!r}")
-    elif _is_cupy_array(x):
+    elif is_cupy_array(x):
         # cupy does not yet have to_device
         return _cupy_to_device(x, device, stream=stream)
-    elif _is_torch_array(x):
+    elif is_torch_array(x):
         return _torch_to_device(x, device, stream=stream)
-    elif _is_dask_array(x):
+    elif is_dask_array(x):
         if stream is not None:
             raise ValueError("The stream argument to to_device() is not supported")
         # TODO: What if our array is on the GPU already?
         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):
@@ -253,4 +288,6 @@ def size(x):
         return None
     return math.prod(x.shape)
 
-__all__ = ['is_array_api_obj', 'array_namespace', 'get_namespace', 'device', 'to_device', 'size']
+__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_jax_array']

tests/test_helpers.py

@@ -0,0 +1,26 @@
+from array_api_compat import (is_numpy_array, is_cupy_array, is_torch_array,
+                              is_dask_array, is_jax_array, is_array_api_obj)
+
+from ._helpers import import_
+
+import pytest
+
+is_functions = {
+    'numpy': 'is_numpy_array',
+    '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())
+@pytest.mark.parametrize('func', is_functions.values())
+def test_is_xp_array(library, func):
+    lib = import_(library)
+    is_func = globals()[func]
+
+    x = lib.asarray([1, 2, 3])
+
+    assert is_func(x) == (func == is_functions[library])
+
+    assert is_array_api_obj(x)

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