EMA
------------------------

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

.. autoclass:: EMA
   :members: __init__, update