jax.experimental 模块

jax.experimental.optix 已经被移动到它自己的 Python 包中 (deepmind/optax)。

jax.experimental.ann 已被移入 jax.lax。

实验性模块

• jax.experimental.array_api 模块
• jax.experimental.checkify 模块
• jax.experimental.host_callback 模块
• jax.experimental.pjit 模块
• jax.experimental.sparse 模块
• jax.experimental.jet 模块
• jax.experimental.custom_partitioning 模块
• jax.experimental.multihost_utils 模块
• jax.experimental.compilation_cache 模块
• jax.experimental.key_reuse 模块
• jax.experimental.mesh_utils 模块
• jax.experimental.serialize_executable 模块
• jax.experimental.shard_map 模块
• jax.experimental.pallas 模块

实验性API

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