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

.. autosummary::
  :toctree: _autosummary

  vmap
  scan
  jit
  remat
  remat_scan
  map_variables
  jvp
  vjp
  grad
  value_and_grad
  custom_vjp
  while_loop
  cond
  switch