Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein
This provides a practical tool for researchers and practitioners in machine learning and statistics to efficiently compute Wasserstein distances and related metrics, though it is incremental as it builds on existing JAX capabilities.
The authors introduced OTT-JAX, a Python toolbox built on JAX for solving optimal transport problems, covering tasks like regularized OT, barycenters, and Gromov-Wasserstein, with features like automatic differentiation and accelerator support.
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}