ott.neural.networks.potentials.PotentialTrainState.tx# PotentialTrainState.tx: optax.GradientTransformation#