OTT-JAX puts optimal transport on the accelerator with Sinkhorn

Optimal transport is a matching problem with weights

OTT-JAX is a JAX powered library to solve a wide variety of problems leveraging optimal transport theory, at scale and on accelerators.

The documentation gives a plain definition first. Optimal transport can be loosely described as the branch of mathematics and optimization that studies matching problems: given two families of points, and a cost function on pairs of points, find a good, low cost way to associate bijectively to every point in the first family another in the second.

The framing matters because it positions the field between two things people already know. The exact algorithm for it is the Hungarian algorithm, which optimally matches two sets of n points using a pairwise cost, but it costs an order of O(n^3) operations and lacks flexibility, since one may want to couple families of different sizes.

Optimal transport extends all of this through faster algorithms, in n^2 or even linear in n, along with generalizations that handle weighted sets of different size, partial matchings and more evolved quadratic matching problems.

So the library is doing two things at once. It replaces the cubic exact solver with an iterative one, and it opens up problem classes that the exact formulation cannot express at all, such as matching where the cost itself is defined between pairs of structures rather than between individual points.

The team behind it is a team of researchers at Apple, with past contributions from Google and Meta researchers, as well as academic partners including TU München, Oxford, ENSAE/IP Paris, ENS Paris and the Hebrew University. The licence is Apache-2.0, and the project is classified as Development Status 4 - Beta.

The Sinkhorn solver is the core, and the geometry is pluggable

The library implements various discrete solvers to match two point clouds, notably the Sinkhorn algorithm, implemented to work on various geometric domains and sped up using various tweaks and extensions.

The tweaks are listed by name: scheduling, momentum, acceleration and initializations. The extension is described as low-rank, which is the standard way of making a Sinkhorn-like iteration scale when the problem is not well conditioned.

What makes this a library rather than a script is the separation of concerns visible in the example. The geometry is a separate object from the solver and the problem. The example imports `pointcloud` from ott.geometry, `linear_problem` from ott.problems.linear, `linear` from ott.solvers and `plot` from ott.tools, then instantiates a PointCloud, and hands it to a solver:

import jax
import jax.numpy as jnp

from ott.geometry import pointcloud
from ott.problems.linear import linear_problem
from ott.solvers import linear
from ott.tools import plot

# sample two point clouds and their weights.
rngs = jax.random.split(jax.random.key(42), 4)
n, m, d = 6, 11, 2
x = jax.random.uniform(rngs[0], (n,d))
y = jax.random.uniform(rngs[1], (m,d))
a = jax.random.uniform(rngs[2], (n,)) +.2
b = jax.random.uniform(rngs[3], (m,)) +.2
a, b = a / jnp.sum(a), b / jnp.sum(b)
# instantiate geometry object to compare point clouds.
geom = pointcloud.PointCloud(x, y)
# compute coupling using the Sinkhorn algorithm.
out = jax.jit(linear.solve)(geom,a,b)

# plot
plot.Plot()(out)

Note the shape of the call. `jax.jit(linear.solve)(geom, a, b)` takes the geometry and the two weight vectors and returns the solution object. The point clouds have different sizes, 6 and 11, which is exactly the flexibility the documentation says the Hungarian algorithm lacks.

The result contains a transport matrix that quantifies the association strength between each point of the first point cloud and one or more points from the second, and the plot tool displays it. The documentation also promises more flexibility to define custom cost functions, objectives and solvers.

Install is pip or conda, and the dependency set tells you what it leans on

OTT-JAX is installed from PyPI:

pip install ott-jax

or with conda via conda-forge:

conda install -c conda-forge ott-jax

The project metadata gives the runtime requirements, and they are a short and specific list. jax at 0.8.0 or newer, jaxopt at 0.8 or newer, lineax at 0.0.7 or newer, numpy at 1.20.0 or newer, and optax at 0.2.4 or newer.

The pairing of jaxopt and optax is the interesting one. jaxopt is the optimiser library that provides the fixed-point iteration style solvers that Sinkhorn is, and optax is the gradient and optimiser library that handles the neural network approaches. lineax is for solving linear systems, which is what a low-rank Sinkhorn variant needs. So the three accelerator-side dependencies correspond to the three things the library does: iterative optimisation, differentiable optimisation, and linear algebra.

The optional extras extend that. A neural extra pulls in flax at 0.11.0 or newer and diffrax at 0.4.1 or newer, which is what the neural network approaches to optimal transport maps are built on. There is a dev extra for pre-commit and a test extra.

The Python requirement is 3.11 or newer, with classifiers for 3.11 through 3.14, and the operating system classifiers cover POSIX Linux, MacOS X and Windows. The package is typed, as the Typing :: Typed classifier records.

The version is dynamic, taken from setuptools-scm, which means it is derived from the git tags rather than written in the file. The package builds with setuptools at 64 or newer.

The harder problems are built on the linear solver

The discrete solvers are not the end of the library. The documentation says these algorithms power the resolution of more advanced problems: Gromov-Wasserstein and Wasserstein barycenter, to compare point clouds in versatile settings.

Gromov-Wasserstein is the one that justifies the whole framework. In a linear optimal transport problem you have two point clouds and a cost between a point in the first and a point in the second, typically a distance. In the quadratic or Gromov-Wasserstein setting you instead have two point clouds and no cross-cost at all, only the internal structure of each cloud, and the task is to align the two structures with each other. That is a much harder optimisation, and it is why the generalisation the README mentions matters: weighted sets of different size, partial matchings, and quadratic matching problems are not variants of the same solver, they are different solvers.

A Wasserstein barycenter is the other one. Given several distributions, it computes the distribution that minimises the average transport cost to all of them, which is a practical way to merge several embeddings into a shared one.

On top of the discrete solvers, the project also proposes implementations of neural network approaches. Given a source and target pair of measure, they output a neural network that seeks to approximate their optimal transport map.

That last piece is the one that changes the character of the library. The discrete solvers give you a coupling matrix. The neural approaches give you a function, which can be evaluated at new points, and the whole thing sits inside JAX so it is differentiable end to end. The keyword list records the intent with implicit differentiation alongside automatic-differentiation and jax.

A young package, still moving, with a research citation attached

The project is pre-1.0 and the release history shows regular small increments rather than a settled API. The recent versions are 0.6.0 from 2025-11-04, 0.5.2 from 2025-10-01 and 0.5.1 from 2025-07-28, and the last push was on 2026-09-18.

The version scheme has moved through 0.5 and 0.6 in a year, so a pin on an old minor version will age.

The citation is a good indicator of how the project sees itself. The reference is Cuturi, Meng-Papaxanthos, Tian, Bunne, Davis and Teboul, Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein, arXiv preprint 2201.12324. The title's phrasing is telling: a toolbox for all things Wasserstein. This is infrastructure for a research community rather than a product with a support commitment.

The documentation is on Read the Docs, the tests run in GitHub Actions with a tests workflow, and the coverage is tracked on Codecov with a .codecov.yml in the repository. The repository layout is conventional for a modern Python package: src/ for the code, tests/ for the suite, docs/ for the documentation, plus a .pre-commit-config.yaml, a .readthedocs.yaml, a CITATION.bib file for the citation, a MANIFEST.in, a code of conduct and a contributing guide.

One entry in the documentation is worth noting for anyone working in biology. The moscot package, for optimal transport analysis of multi-omics data, uses OTT as a backbone. That is a concrete signal that the library is used as a component by other research tools, rather than only by its own authors.

Where this sits against other optimal transport libraries

The comparison worth making is with the other two common approaches to optimal transport in Python.

The first is POT, the Python Optimal Transport library. It is the long-standing general-purpose package, with a broad set of solvers and estimators, and it is not accelerator-oriented. It runs on CPU with numpy and scipy, and it solves a lot of problems that do not need a GPU. For a dataset that fits in memory and a single machine, that is often the right tool, and it is the one most likely to have an example that matches your exact problem already written.

The second is geomloss, which took the approach of implementing the Sinkhorn iteration with custom autodiff rules so that gradients through the transport plan were cheap. That was a significant idea and it is why optimal transport appeared in so many differentiable pipelines. OTT-JAX's positioning is the natural successor to that: JAX's autodiff and jit mean the compiler handles the differentiation and the fusion of the iteration, without a hand-written backward pass per solver.

The third is the neural approach, and here OTT differs from geomloss in ambition. Instead of differentiating through a discrete solver, OTT also implements neural network approaches that directly output an approximation of the optimal transport map, which is a different computational object.

The practical comparison is therefore about scale and about which problem you have. If you want exact assignment on a small problem, use the Hungarian algorithm and forget about libraries. If you want a large linear problem and you already use JAX for the rest of your model, OTT-JAX fits. If your problem is the quadratic kind, where the cost is between structures rather than points, OTT is one of the few libraries that exposes it directly.

Editorial conclusion

Adopt OTT-JAX if your problem is a matching between two point clouds at a scale where the Hungarian algorithm is out of reach, or if you need the harder quadratic variants like Gromov-Wasserstein that general matching libraries do not expose. Skip it if your problem is small enough for exact assignment, since the iterative solver trades exactness for speed, or if you need a non-Python implementation, because this is a JAX library and inherits JAX's execution model. Verify first that your accelerator and JAX version line up with the declared jax>=0.8.0 requirement, that you can accept a Beta status package with a dynamic version, and that the geometry you need has a matching cost implementation.

Frequently asked questions

Official sources

1. License: Apache-2.0
2. ott-jax/ott on GitHub
3. Project website
4. README
5. Releases