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

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

.. flax_module::
  :module: flax.nnx
  :class: BatchNorm

.. flax_module::
  :module: flax.nnx
  :class: LayerNorm

.. flax_module::
  :module: flax.nnx
  :class: RMSNorm