ott.neural.solvers.map_estimator.MapEstimator.setup