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.pjit 模块
• jax.experimental.serialize_executable 模块
• jax.experimental.shard_map 模块
• jax.experimental.sparse 模块

实验性 API

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