jax.experimental 模块

jax.experimental.optix 已被移至其独立的 Python 包（deepmind/optax）中。

jax.experimental.ann 已被移至 jax.lax 中。

实验性模块

• jax.experimental.checkify 模块
• jax.experimental.compilation_cache 模块
• jax.experimental.custom_partitioning 模块
• jax.experimental.jet 模块
• jax.experimental.key_reuse 模块
• jax.experimental.mesh_utils 模块
• jax.experimental.multihost_utils 模块
• jax.experimental.pallas 模块
• jax.experimental.random 模块
• jax.experimental.serialize_executable 模块
• jax.experimental.sparse 模块