jax.experimental.array_api module

This module includes experimental JAX support for the Python array API standard. Support for this is currently experimental and not fully complete.

Example Usage:

>>> from jax.experimental import array_api as xp

>>> xp.__array_api_version__
'2022.12'

>>> arr = xp.arange(1000)

>>> arr.sum()
Array(499500, dtype=int32)

The xp namespace is the array API compliant analog of jax.numpy, and implements most of the API listed in the standard.