jax.experimental 模块

jax.experimental.optix 已移至其自己的 Python 包 (deepmind/optax)。

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

实验性模块

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

实验性 API

enable_x64([new_val])	用于临时启用 X64 模式的实验性上下文管理器。
disable_x64()	用于临时禁用 X64 模式的实验性上下文管理器。