Add jax.experimental.array_api interface.

.github/workflows/jax-array-api.yml

@@ -0,0 +1,48 @@
+name: JAX Array API
+
+on:
+  workflow_dispatch: # allows triggering the workflow run manually
+  pull_request:  # Automatically trigger on pull requests affecting this file
+    branches:
+      - main
+    paths:
+      - '**workflows/jax-array-api.yml'
+      - '**experimental/array_api/**'
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.11]
+
+    steps:
+    - name: Checkout jax
+      uses: actions/checkout@v3
+    - name: Checkout array-api-tests
+      uses: actions/checkout@v3
+      with:
+        repository: data-apis/array-api-tests
+        submodules: 'true'
+        path: 'array-api-tests'
+    - name: Fix array-apis bug
+      # Temporary workaround for https://github.com/data-apis/array-api/issues/631 
+      run: |
+        sed -i -e 's/\\/\\\\/g' array-api-tests/array-api/spec/API_specification/signatures/*.py
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v1
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install .[cpu]
+        python -m pip install hypothesis!=6.88.4  # 6.88.4 leads to a strange error
+        python -m pip install -r array-api-tests/requirements.txt
+    - name: Run the test suite
+      env:
+        ARRAY_API_TESTS_MODULE: jax.experimental.array_api
+        JAX_ENABLE_X64: 'true'
+      run: |
+        cd ${GITHUB_WORKSPACE}/array-api-tests
+        pytest --ci array_api_tests --max-examples=5  --skips-file ${GITHUB_WORKSPACE}/array-api-skips.txt

array-api-skips.txt

@@ -0,0 +1,43 @@
+# Known failures for the array api tests.
+
+# JAX doesn't yet support scalar boolean indexing
+array_api_tests/test_array_object.py::test_getitem_masking
+
+# Hypothesis warning
+array_api_tests/meta/test_hypothesis_helpers.py::test_symmetric_matrices
+
+# Test suite attempts in-place mutation:
+array_api_tests/test_special_cases.py::test_binary
+array_api_tests/test_special_cases.py::test_iop
+array_api_tests/test_special_cases.py::test_nan_propagation
+array_api_tests/test_special_cases.py::test_unary
+array_api_tests/test_array_object.py::test_setitem
+array_api_tests/test_creation_functions.py::test_asarray_arrays
+array_api_tests/test_linalg.py::test_matrix_power
+array_api_tests/test_linalg.py::test_solve
+
+# Overflow errors due to hypothesis generating integers that overflow int64
+array_api_tests/test_operators_and_elementwise_functions.py::test_multiply[__mul__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_square
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_and[__and__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_and[__iand__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_or[__or__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_or[__ior__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__lshift__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__ilshift__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__rshift__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__irshift__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_xor[__xor__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_xor[__ixor__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[__floordiv__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[__ifloordiv__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_greater_equal[__ge__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_less[__lt__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_not_equal[__ne__(x, s)]
+
+# JAX's NaN sorting doesn't match specification
+array_api_tests/test_set_functions.py::test_unique_all
+array_api_tests/test_set_functions.py::test_unique_counts
+array_api_tests/test_set_functions.py::test_unique_inverse
+array_api_tests/test_set_functions.py::test_unique_values
+array_api_tests/test_sorting_functions.py::test_argsort

jax/experimental/array_api/__init__.py

@@ -0,0 +1,195 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from jax.experimental.array_api._version import __array_api_version__ as __array_api_version__
+
+from jax.experimental.array_api import linalg as linalg
+
+from jax.experimental.array_api._constants import (
+    e as e,
+    inf as inf,
+    nan as nan,
+    newaxis as newaxis,
+    pi as pi,
+)
+
+from jax.experimental.array_api._creation_functions import (
+    arange as arange,
+    asarray as asarray,
+    empty as empty,
+    empty_like as empty_like,
+    eye as eye,
+    from_dlpack as from_dlpack,
+    full as full,
+    full_like as full_like,
+    linspace as linspace,
+    meshgrid as meshgrid,
+    ones as ones,
+    ones_like as ones_like,
+    tril as tril,
+    triu as triu,
+    zeros as zeros,
+    zeros_like as zeros_like,
+)
+
+from jax.experimental.array_api._data_type_functions import (
+    astype as astype,
+    can_cast as can_cast,
+    finfo as finfo,
+    iinfo as iinfo,
+    isdtype as isdtype,
+    result_type as result_type,
+)
+
+from jax.experimental.array_api._dtypes import (
+    bool as bool,
+    int8 as int8,
+    int16 as int16,
+    int32 as int32,
+    int64 as int64,
+    uint8 as uint8,
+    uint16 as uint16,
+    uint32 as uint32,
+    uint64 as uint64,
+    float32 as float32,
+    float64 as float64,
+    complex64 as complex64,
+    complex128 as complex128,
+)
+
+from jax.experimental.array_api._elementwise_functions import (
+    abs as abs,
+    acos as acos,
+    acosh as acosh,
+    add as add,
+    asin as asin,
+    asinh as asinh,
+    atan as atan,
+    atan2 as atan2,
+    atanh as atanh,
+    bitwise_and as bitwise_and,
+    bitwise_invert as bitwise_invert,
+    bitwise_left_shift as bitwise_left_shift,
+    bitwise_or as bitwise_or,
+    bitwise_right_shift as bitwise_right_shift,
+    bitwise_xor as bitwise_xor,
+    ceil as ceil,
+    conj as conj,
+    cos as cos,
+    cosh as cosh,
+    divide as divide,
+    equal as equal,
+    exp as exp,
+    expm1 as expm1,
+    floor as floor,
+    floor_divide as floor_divide,
+    greater as greater,
+    greater_equal as greater_equal,
+    imag as imag,
+    isfinite as isfinite,
+    isinf as isinf,
+    isnan as isnan,
+    less as less,
+    less_equal as less_equal,
+    log as log,
+    log10 as log10,
+    log1p as log1p,
+    log2 as log2,
+    logaddexp as logaddexp,
+    logical_and as logical_and,
+    logical_not as logical_not,
+    logical_or as logical_or,
+    logical_xor as logical_xor,
+    multiply as multiply,
+    negative as negative,
+    not_equal as not_equal,
+    positive as positive,
+    pow as pow,
+    real as real,
+    remainder as remainder,
+    round as round,
+    sign as sign,
+    sin as sin,
+    sinh as sinh,
+    sqrt as sqrt,
+    square as square,
+    subtract as subtract,
+    tan as tan,
+    tanh as tanh,
+    trunc as trunc,
+)
+
+from jax.experimental.array_api._indexing_functions import (
+    take as take,
+)
+
+from jax.experimental.array_api._manipulation_functions import (
+    broadcast_arrays as broadcast_arrays,
+    broadcast_to as broadcast_to,
+    concat as concat,
+    expand_dims as expand_dims,
+    flip as flip,
+    permute_dims as permute_dims,
+    reshape as reshape,
+    roll as roll,
+    squeeze as squeeze,
+    stack as stack,
+)
+
+from jax.experimental.array_api._searching_functions import (
+    argmax as argmax,
+    argmin as argmin,
+    nonzero as nonzero,
+    where as where,
+)
+
+from jax.experimental.array_api._set_functions import (
+    unique_all as unique_all,
+    unique_counts as unique_counts,
+    unique_inverse as unique_inverse,
+    unique_values as unique_values,
+)
+
+from jax.experimental.array_api._sorting_functions import (
+    argsort as argsort,
+    sort as sort,
+)
+
+from jax.experimental.array_api._statistical_functions import (
+    max as max,
+    mean as mean,
+    min as min,
+    prod as prod,
+    std as std,
+    sum as sum,
+    var as var
+)
+
+from jax.experimental.array_api._utility_functions import (
+    all as all,
+    any as any,
+)
+
+from jax.experimental.array_api._linear_algebra_functions import (
+    matmul as matmul,
+    matrix_transpose as matrix_transpose,
+    tensordot as tensordot,
+    vecdot as vecdot,
+)
+
+from jax.experimental.array_api import _array_methods
+_array_methods.add_array_object_methods()
+del _array_methods

jax/experimental/array_api/_array_methods.py

@@ -0,0 +1,45 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any, Callable, Optional, Union
+
+import jax
+from jax._src.array import ArrayImpl
+from jax.experimental.array_api._version import __array_api_version__
+
+from jax._src.lib import xla_extension as xe
+
+
+def _array_namespace(self, /, *, api_version: None | str = None):
+  if api_version is not None and api_version != __array_api_version__:
+    raise ValueError(f"{api_version=!r} is not available; "
+                     f"available versions are: {[__array_api_version__]}")
+  return jax.experimental.array_api
+
+
+def _to_device(self, device: xe.Device | Callable[[], xe.Device], /, *,
+               stream: Optional[Union[int, Any]] = None):
+  if stream is not None:
+    raise NotImplementedError("stream argument of array.to_device()")
+  # The type of device is defined by Array.device. In JAX, this is a callable that
+  # returns a device, so we must handle this case to satisfy the API spec.
+  return jax.device_put(self, device() if callable(device) else device)
+
+
+def add_array_object_methods():
+  # TODO(jakevdp): set on tracers as well?
+  setattr(ArrayImpl, "__array_namespace__", _array_namespace)
+  setattr(ArrayImpl, "to_device", _to_device)

jax/experimental/array_api/_constants.py

@@ -0,0 +1,7 @@
+import numpy as np
+
+e = np.e
+inf = np.inf
+nan = np.nan
+newaxis = np.newaxis
+pi = np.pi