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.