Normalization
------------------------

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

.. autoclass:: BatchNorm
   :members:
.. autoclass:: LayerNorm
   :members:
.. autoclass:: RMSNorm
   :members: