---
id: ott-jax
version: "0.6.0"
license: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record)
license_treatment: permissive
maintenance: active
---

# ott-jax — Optimal Transport Tools in JAX

License: permissive · Maintenance: active · Downloads: 94.1K/mo

## 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 above — 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

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

- License: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record) (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 94.1K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

optimal transport matching, wasserstein distance computation, sinkhorn algorithm jax, point cloud alignment, gromov wasserstein, transport map neural network, jax accelerated optimization, optimal-transport, jax-accelerated, point-cloud-matching

[View on SkillFed](https://skillfed.io/packages/ott-jax) · [View on PyPI](https://pypi.org/project/ott-jax/)