Optimal Transport Tools (OTT) documentation

Code hosted on Github. To install, clone that repo or simply run pip install ott-jax.


OTT is a JAX package that bundles a few utilities to compute and differentiate the solution to optimal transport problems. OTT can help you compute Wasserstein distances between weighted clouds of points (or histograms), using a cost (e.g. a distance) between individual points.

To that end OTT uses a sturdy and versatile implementation of the Sinkhorn algorithm 1 2. This implementation takes advantage of several JAX features, such as Just-in-time (JIT) compilation, auto-vectorization (VMAP), and both automatic and/or implicit differentiation. Basic tutorial snippets are provided below, along with a more advanced example to process single-cell genomics data 3.


There are currently three packages, geometry, core and tools, playing the following roles:

  • geometry defines classes that describe two point clouds paired with a cost function (simpler geometries are also implemented, such as that defined by points supported on a multi-dimensional grids with a separable cost 4). A geometry, along with weight vectors a and b, describes an OT problem. core algorithms can solve such problems, using subroutines provided by geometries;

  • core contains the Sinkhorn algorithm, the main workhorse to solve OT in this package, as well as variants that can be used to compute barycenters of several measures;

  • tools provides an interface to exploit OT solutions, as produced by core functions. Such tasks include instantiating OT matrices, computing approximations to Wasserstein distances 5 6, or computing differentiable sort and quantile operations 7.

Indices and tables


  1. Cuturi, Sinkhorn Distances: Lightspeed Computation of Optimal Transport, NIPS’13.

  1. Peyré, M. Cuturi, Computational Optimal Transport, FNT in ML, 2019.

  1. Schiebinger et al., Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming, Cell 176, 928–943.

  1. Solomon et al, Convolutional Wasserstein distances: efficient optimal transportation on geometric domains, ACM ToG, SIGGRAPH’15.

  1. Genevay et al., Learning Generative Models with Sinkhorn Divergences, AISTATS’18.

  1. Séjourné et al., Sinkhorn Divergences for Unbalanced Optimal Transport, arXiv:1910.12958.

  1. Cuturi et al. Differentiable Ranking and Sorting using Optimal Transport, NeurIPS’19.