jax.experimental.rst

jax.experimental package

.. toctree::
    :maxdepth: 1

    jax.experimental.host_callback
    jax.experimental.loops
    jax.experimental.optimizers
    jax.experimental.optix
    jax.experimental.stax

.. automodule:: jax.experimental