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 模式。