ott.core package#

OTT core libraries: the engines behind most computations happening in OTT.

The core package contains definitions of various OT problems, starting from the most simple, the linear OT problem, to more advanced problems such as quadratic, or involving multiple measures, the barycenter problem. We follow with the classic sinkhorn routine (essentially a wrapper for the Sinkhorn solver class) 1, 2. We also provide an analogous low-rank Sinkhorn solver 3 to handle very large instances. Both are used within our Wasserstein barycenter solvers 4, 5 as well as our Gromov-Wasserstein solver 6, 7. We also provide an implementation of input convex neural networks 8, a NN that can be used to estimate OT 9.

OT Problems#

problems.LinearProblem(geom[, a, b, tau_a, ...])

Holds the definition of a linear regularized OT problem and some tools.

quad_problems.QuadraticProblem(geom_xx, geom_yy)

Holds the definition of the quadratic regularized OT problem.

bar_problems.BarycenterProblem([y, b, ...])

Holds the definition of a linear regularized OT problem and some tools.

Sinkhorn#

sinkhorn.sinkhorn(geom[, a, b, tau_a, ...])

A Jax version of Sinkhorn's algorithm.

sinkhorn.Sinkhorn([lse_mode, threshold, ...])

A Sinkhorn solver for linear reg-OT problem implemented as a pytree.

Low-Rank Sinkhorn#

sinkhorn_lr.LRSinkhorn([rank, gamma, ...])

A Low-Rank Sinkhorn solver for linear reg-OT problems.

Barycenters (Entropic and LR)#

discrete_barycenter.discrete_barycenter(geom, a)

Compute discrete barycenter using https://arxiv.org/abs/2006.02575.

Gromov-Wasserstein (Entropic and LR)#

gromov_wasserstein.gromov_wasserstein(...[, ...])

Wrapper to solve a Gromov Wasserstein problem.

Neural Potentials#

icnn.ICNN(dim_hidden[, init_std, init_fn, ...])

Input convex neural network (ICNN) architeture.

References#

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

2
  1. Séjourné, Sinkhorn Divergences for Unbalanced Optimal Transport , NeurIPS 2019.

3
  1. Scetbon et al., Low-Rank Sinkhorn Factorization , ICML 2021.

4

J.D. Benamou et al., Iterative Bregman Projections for Regularized Transportation Problems , SIAM J. Sci. Comput. 37(2), A1111-A1138.

5
  1. Janati et al., Debiased Sinkhorn Barycenters , ICML 2020.

6
  1. Memoli, Gromov–Wasserstein distances and the metric approach to object matching , FOCM 2011.

7
  1. Scetbon et al., Linear-Time Gromov Wasserstein Distances using Low Rank Couplings and Costs, Arxiv.

8
  1. Amos, L. Xu, J. Z. Kolter, Input Convex Neural Networks, ICML 2017.

9

Ashok Vardhan Makkuva, Amirhossein Taghvaei, Sewoong Oh, Jason D. Lee, Optimal transport mapping via input convex neural networks , ICML 2020