Transformations
----------------------

.. automodule:: flax.linen.transforms
.. currentmodule:: flax.linen

.. autofunction:: vmap
.. autofunction:: scan
.. autofunction:: jit
.. autofunction:: remat
.. autofunction:: remat_scan
.. autofunction:: map_variables
.. autofunction:: jvp
.. autofunction:: vjp
.. autofunction:: custom_vjp
.. autofunction:: while_loop
.. autofunction:: cond
.. autofunction:: switch