jax.experimental module

jax.experimental.optix has been moved into its own Python package (deepmind/optax).

jax.experimental.ann has been moved into jax.lax.

Experimental Modules

• jax.experimental.checkify module
• jax.experimental.compilation_cache module
• jax.experimental.custom_dce module
• jax.experimental.custom_partitioning module
• jax.experimental.jet module
• jax.experimental.key_reuse module
• jax.experimental.mesh_utils module
• jax.experimental.multihost_utils module
• jax.experimental.pallas module
• jax.experimental.pjit module
• jax.experimental.serialize_executable module
• jax.experimental.shard_map module
• jax.experimental.sparse module

Experimental APIs

enable_x64([new_val])	Experimental context manager to temporarily enable X64 mode.
disable_x64()	Experimental context manager to temporarily disable X64 mode.