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