ott.neural.networks.potentials.PotentialTrainState
==================================================

.. currentmodule:: ott.neural.networks.potentials

.. autoclass:: PotentialTrainState
    
    .. rubric:: Methods

    .. autosummary::
        :toctree: .
    
        ~PotentialTrainState.apply_gradients
        ~PotentialTrainState.create
        ~PotentialTrainState.replace
    
    .. rubric:: Attributes

    .. autosummary::
        :toctree: .
    
        ~PotentialTrainState.potential_value_fn
        ~PotentialTrainState.potential_gradient_fn
        ~PotentialTrainState.step
        ~PotentialTrainState.apply_fn
        ~PotentialTrainState.params
        ~PotentialTrainState.tx
        ~PotentialTrainState.opt_state