jax.experimental.array_api 模块

备注

自 JAX v0.4.32 起，jax.experimental.array_api 模块已被弃用，导入 jax.experimental.array_api 不再必要。{mod}`jax.numpy` 默认直接实现了数组 API 标准。详情请参见 Python 数组 API 标准。

此模块包括对 Python 数组 API 标准 的实验性 JAX 支持。目前对此的支持是实验性的，并未完全完成。

示例用法:

>>> from jax.experimental import array_api as xp

>>> xp.__array_api_version__
'2023.12'

>>> arr = xp.arange(1000)

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

xp 命名空间是符合数组API标准的 jax.numpy 的模拟，并实现了标准中列出的大部分API。