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_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.random module
• jax.experimental.serialize_executable module
• jax.experimental.sparse module