transforms
------------------------

.. automodule:: flax.nnx
.. currentmodule:: flax.nnx

.. autoclass:: Jit
   :members:
.. autoclass:: Remat
   :members:
.. autoclass:: Scan
   :members:
.. autoclass:: Vmap
   :members:

.. autofunction:: grad
.. autofunction:: jit
.. autofunction:: remat
.. autofunction:: scan
.. autofunction:: value_and_grad
.. autofunction:: vmap
.. autofunction:: eval_shape
.. autofunction:: custom_vjp
.. autofunction:: cond
.. autofunction:: switch
.. autofunction:: while_loop
.. autofunction:: fori_loop
