ott-jax

What it is and what it does

OTT-JAX is a JAX-powered library for solving optimal transport problems at scale on accelerators. It implements discrete solvers—notably the Sinkhorn algorithm with various optimizations (scheduling, momentum, low-rank approximations)—to compute couplings between point clouds. Beyond basic matching, it supports advanced problems like Gromov-Wasserstein distance (comparing point clouds in different geometric spaces) and Wasserstein barycenters (finding representative point sets). The library also includes neural network approaches that learn transport maps between measure pairs, enabling both exact and approximate solutions depending on problem size and accuracy needs.

The package is built on six core dependencies: jax for autodifferentiation and GPU/TPU acceleration, jaxopt and optax for optimization, lineax for linear algebra, numpy for array operations, and typing_extensions for type hints. It supports Python 3.9 through 3.13 and runs on Linux, macOS, and Windows. The codebase is maintained by researchers at Apple with contributions from Google, Meta, and academic partners, reflecting active development and research backing.

Use it for:

• Compute the optimal coupling matrix between two point clouds to measure similarity or perform alignment.
• Compare distributions in different geometric spaces using Gromov-Wasserstein distance for domain adaptation or shape matching.
• Train a neural network to approximate the optimal transport map between two measures for fast inference on new data.
• Compute Wasserstein barycenters to find representative point sets that minimize transport cost to a collection of input distributions.
• Integrate optimal transport into machine learning pipelines via JAX's autodiff for end-to-end gradient-based optimization.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

OTT-JAX solves optimal transport problems—matching and comparing point clouds—using JAX-accelerated algorithms like Sinkhorn, including support for Gromov-Wasserstein and Wasserstein barycenter computations.

Yes. OTT-JAX is actively maintained, has low install friction, carries a permissive Apache 2.0 license, and implements a mature, well-documented suite of optimal transport algorithms. It is suitable for research, production machine learning, and scientific computing where point cloud matching, distribution comparison, or transport-based optimization is needed. No known security vulnerabilities.

Install

ott-jax on PyPI

pip install ott-jax

uv add ott-jax

poetry add ott-jax

Installing ott-jax

Before you install

Low install friction; pure Python wheel with six runtime dependencies (jax, jaxopt, lineax, numpy, typing_extensions, optax). Actively maintained with last commit in June 2026 and 752 GitHub stars.

License in practice

Apache License 2.0 (permissive): you may use, modify, and distribute freely in commercial or private projects, provided you include a copy of the license and document any changes to the source.

Quickstart

pip install ott-jax

import jax
import jax.numpy as jnp
from ott.geometry import pointcloud
from ott.solvers import linear

x = jax.random.uniform(jax.random.key(0), (10, 2))
y = jax.random.uniform(jax.random.key(1), (10, 2))
geom = pointcloud.PointCloud(x, y)
out = jax.jit(linear.solve)(geom)

Requires JAX and its dependencies (including a compatible NumPy); compute-heavy operations benefit from GPU/TPU access but run on CPU.

Verify before relying

• Whether neural network transport map training requires additional setup or hyperparameter tuning beyond the documented API.
• Performance scaling characteristics on very large point clouds (dimensionality and cardinality thresholds).
• Compatibility with JAX's latest versions and any breaking changes in upstream dependencies.

Package facts

Evidence: ott_jax-0.6.0-py3-none-any.whl

Keywords: optimal transport, gromov wasserstein, sinkhorn, low-rank sinkhorn, sinkhorn divergences, wasserstein, wasserstein barycenter, jax, autodiff, implicit differentiation

Tags

More Mathematics packages