arXiv · 2201.12324
Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein
Abstract
Optimal transport tools (OTT-JAX) is a Python toolbox that can solve optimal transport problems between point clouds and histograms. The toolbox builds on various JAX features, such as automatic and custom reverse mode differentiation, vectorization, just-in-time compilation and accelerators support. The toolbox covers elementary computations, such as the resolution of the regularized OT problem, and more advanced extensions, such as barycenters, Gromov-Wasserstein, low-rank solvers, estimation of convex maps, differentiable generalizations of quantiles and ranks, and approximate OT between Gaussian mixtures. The toolbox code is available at \texttt{https://github.com/ott-jax/ott}
Explore related subjects
Keep this discovery
Marco Cuturi, Laetitia Meng-Papaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, Olivier Teboul. 2022-01-28. Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein. https://arxiv.org/abs/2201.12324
Cite the original work for its findings. Save a collection to share your selection of sources.