Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein
Machine Learning
2022-01-31 v1 Machine Learning
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}
Keywords
Cite
@article{arxiv.2201.12324,
title = {Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein},
author = {Marco Cuturi and Laetitia Meng-Papaxanthos and Yingtao Tian and Charlotte Bunne and Geoff Davis and Olivier Teboul},
journal= {arXiv preprint arXiv:2201.12324},
year = {2022}
}
Comments
4 pages