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:: cond
