nn
----------------------------

Neural network layers and activation functions used in NNX :class:`Module`'s.
See the `NNX page <https://flax.readthedocs.io/en/latest/nnx/index.html>`__ for more details.

.. toctree::
  :maxdepth: 3

  activations
  attention
  dtypes
  initializers
  linear
  lora
  normalization
  pooling
  recurrent
  stochastic