flax.training package
=====================

Train state
------------------------

.. currentmodule:: flax.training.train_state

.. autoclass:: TrainState
    :members: apply_gradients, create