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

.. automodule:: flax.nnx
.. currentmodule:: flax.nnx
.. autofunction:: grad
.. autofunction:: jit
.. autofunction:: shard_map
.. autofunction:: remat
.. autofunction:: scan
.. autoclass:: Carry
.. autofunction:: value_and_grad
.. autofunction:: vmap
.. autofunction:: eval_shape
.. autofunction:: custom_vjp
.. autofunction:: vjp
.. autofunction:: jvp
.. autofunction:: cond
.. autofunction:: switch
.. autofunction:: while_loop
.. autofunction:: fori_loop