Warning

This page was created from a pull request.

jax.experimental package

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

• jax.experimental.host_callback module
• jax.experimental.loops module
• jax.experimental.optimizers module
• jax.experimental.stax module