flax.nnx
------------------------

Experimental API. See the `NNX page <https://flax.readthedocs.io/en/latest/nnx/index.html>`__ for more details.

.. toctree::
  :maxdepth: 3

  graph
  compat
  object
  module
  nn/index
  rnglib
  spmd
  state
  training/index
  transforms
  variables
  helpers
  visualization
  filterlib
  bridge