-
Notifications
You must be signed in to change notification settings - Fork 2.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add jax.experimental.array_api interface.
- Loading branch information
Showing
20 changed files
with
1,615 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,7 @@ | ||
import numpy as np | ||
|
||
e = np.e | ||
inf = np.inf | ||
nan = np.nan | ||
newaxis = np.newaxis | ||
pi = np.pi |
Oops, something went wrong.