variables
------------------------

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

.. autoclass:: BatchStat
   :members:
.. autoclass:: Cache
   :members:
.. autoclass:: Empty
   :members:
.. autoclass:: Intermediate
   :members:
.. autoclass:: Param
   :members:
.. autoclass:: Variable
   :members:
.. autoclass:: VariableMetadata
   :members:
.. autoclass:: VariableState
   :members:

.. autofunction:: with_metadata