ott.neural.networks.potentials.PotentialTrainState.potential_gradient_fn

PotentialTrainState.potential_gradient_fn: Callable[[FrozenDict[str, Array]], Callable[[Array], Array]]