ott-jax
Optimal Transport Tools in JAX
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires JAX and its dependencies (including a compatible NumPy); compute-heavy operations benefit from GPU/TPU access but run on CPU.
- 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 · maintenance · safety
permissive license (permissive) — 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.
last release 2025-11-04 (283 days) · last repo commit 2026-06-14 · 752 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 94,080 downloads/mo, #13,344 on PyPI
Alternatives
Verify before relying
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)- 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.
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 on it.
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
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.
Requires JAX and its dependencies (including a compatible NumPy); compute-heavy operations benefit from GPU/TPU access but run on CPU.
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)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 | permissive license permissive |
| Python support | Supports the current Python release >=3.9 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 6 packagesjaxjaxoptlineaxnumpytyping_extensionsoptax |
| Maintenance | Actively maintained 283 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 94,080 / month, #13,344 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 4 - BetaIntended Audience :: DevelopersIntended Audience :: Science/ResearchLicense :: OSI Approved :: Apache Software LicenseNatural Language :: EnglishOperating System :: MacOS :: MacOS XOperating System :: Microsoft :: WindowsOperating System :: POSIX :: LinuxProgramming Language :: Python :: 3Programming Language :: Python :: 3.10Programming Language :: Python :: 3.11Programming Language :: Python :: 3.12Programming Language :: Python :: 3.13Programming Language :: Python :: 3.9Topic :: Scientific/Engineering :: MathematicsTyping :: Typed |
Evidence: ott_jax-0.6.0-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “optimal transport matching”
- ott-jaxOTT-JAX solves optimal transport problems—matching and comparing…
- geomlossComputes geometric loss functions for comparing point clouds, images,…
- POTPOT provides solvers for optimal transport problems, including…
Give your agent the search over MCP, or paste the wish link into any chat.
More Mathematics packages
NetworkX provides data structures and algorithms for creating, analyzing, and manipulating graphs and networks, supporting everything from simple undirected graphs to complex directed and weighted networks.
kiwisolver is a Python binding to a fast C++ implementation of the Cassowary constraint solver, enabling you to solve systems of linear constraints and inequalities.
Install it if you need to solve constraint systems; skip it if you only need simple linear algebra.
SymPy is a Python library for symbolic mathematics, performing algebraic manipulation, calculus, equation solving, and mathematical expression simplification without numerical approximation.
ContourPy calculates contours of 2D quadrilateral grids using C++11 algorithms wrapped in Python, offering serial and multithreaded implementations without requiring Matplotlib as a dependency.
PyTorch provides GPU-accelerated tensor computation and automatic differentiation for building and training deep neural networks in Python.
onnxruntime loads and executes Open Neural Network Exchange (ONNX) models with a focus on inference performance across CPUs and accelerators.
Install it if you have ONNX models to run in production or development.
See also geomloss · POT · lineax · optimistix · jax-cuda12-pjrt · jax · jax-cuda13-pjrt · lapx · lap