ott.neural.networks.potentials.PotentialTrainState.opt_state# PotentialTrainState.opt_state: optax.OptState#