# Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein

arXiv: https://arxiv.org/abs/2201.12324 · HTML: https://mlsanity.com/paper/2201.12324/

## Key facts

- Cited by 101 papers in the map (data through 14 Oct 2026).
- Builds on 8 of its 17 references that are in the map.
- First posted on arXiv 28 Jan 2022.

## Builds on (most cited first, 8 of 8)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem](https://mlsanity.com/paper/1905.11882/) (2019)
- [A Wasserstein-type distance in the space of Gaussian Mixture Models](https://mlsanity.com/paper/1907.05254/) (2019)
- [Low-Rank Sinkhorn Factorization](https://mlsanity.com/paper/2103.04737/) (2021)
- [Linear-Time Gromov Wasserstein Distances using Low Rank Couplings and Costs](https://mlsanity.com/paper/2106.01128/) (2021)
- [Differential Properties of Sinkhorn Approximation for Learning with Wasserstein Distance](https://mlsanity.com/paper/1805.11897/) (2018)

## Built on by (most cited first, 30 of 101)

- [Soft Diffusion: Score Matching for General Corruptions](https://mlsanity.com/paper/2209.05442/) (2022)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [Recent Advances in Optimal Transport for Machine Learning](https://mlsanity.com/paper/2306.16156/) (2023)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Sequential Controlled Langevin Diffusions](https://mlsanity.com/paper/2412.07081/) (2024)
- [Optimal Transport for Offline Imitation Learning](https://mlsanity.com/paper/2303.13971/) (2023)
- [Aligned Diffusion Schr"odinger Bridges](https://mlsanity.com/paper/2302.11419/) (2023)
- [Simple ReFlow: Improved Techniques for Fast Flow Models](https://mlsanity.com/paper/2410.07815/) (2024)
- [Debias Coarsely, Sample Conditionally: Statistical Downscaling through Optimal Transport and Probabilistic Diffusion Models](https://mlsanity.com/paper/2305.15618/) (2023)
- [Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation](https://mlsanity.com/paper/2311.15100/) (2023)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [Get more for less: Principled Data Selection for Warming Up Fine-Tuning in LLMs](https://mlsanity.com/paper/2405.02774/) (2024)
- [DySLIM: Dynamics Stable Learning by Invariant Measure for Chaotic Systems](https://mlsanity.com/paper/2402.04467/) (2024)
- [Underdamped Diffusion Bridges with Applications to Sampling](https://mlsanity.com/paper/2503.01006/) (2025)
- [Meta Optimal Transport](https://mlsanity.com/paper/2206.05262/) (2022)
- [GENOT: Entropic (Gromov) Wasserstein Flow Matching with Applications to Single-Cell Genomics](https://mlsanity.com/paper/2310.09254/) (2023)
- [Wasserstein Flow Matching: Generative modeling over families of distributions](https://mlsanity.com/paper/2411.00698/) (2024)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Addressing Misspecification in Simulation-based Inference through Data-driven Calibration](https://mlsanity.com/paper/2405.08719/) (2024)
- [Multi-marginal Schr"odinger Bridges with Iterative Reference Refinement](https://mlsanity.com/paper/2408.06277/) (2024)
- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [Multivariate Conformal Prediction using Optimal Transport](https://mlsanity.com/paper/2502.03609/) (2025)
- [Low-rank Optimal Transport: Approximation, Statistics and Debiasing](https://mlsanity.com/paper/2205.12365/) (2022)
- [Causal Modeling with Stationary Diffusions](https://mlsanity.com/paper/2310.17405/) (2023)
- [Self-Consistent Velocity Matching of Probability Flows](https://mlsanity.com/paper/2301.13737/) (2023)
- [Monge, Bregman and Occam: Interpretable Optimal Transport in High-Dimensions with Feature-Sparse Maps](https://mlsanity.com/paper/2302.04065/) (2023)
- [Rethinking Initialization of the Sinkhorn Algorithm](https://mlsanity.com/paper/2206.07630/) (2022)
- [Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers](https://mlsanity.com/paper/2404.09411/) (2024)
- [Discovering Conservation Laws using Optimal Transport and Manifold Learning](https://mlsanity.com/paper/2208.14995/) (2022)

## Abstract (arXiv)

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 texttthttps://github.com/ott-jax/ott