ott.solvers.linear#

Linear solvers are the bread-and-butter of OT solvers. They can be called on their own, either the Sinkhorn Sinkhorn or Low-Rank LRSinkhorn solvers, to match two datasets. They also appear as subroutines for more advanced solvers in the ott.solvers module, notably ott.solvers.quadratic.

Sinkhorn Solvers#

solve(geom[, a, b, tau_a, tau_b, rank])

Solve linear regularized OT problem using Sinkhorn iterations.

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

Sinkhorn solver.

sinkhorn.SinkhornState([errors, fu, gv, ...])

Holds the state variables used to solve OT with Sinkhorn.

sinkhorn.SinkhornOutput([f, g, errors, ...])

Holds the output of a Sinkhorn solver applied to a problem.

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

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

sinkhorn_lr.LRSinkhornState(q, r, g, gamma, ...)

State of the Low Rank Sinkhorn algorithm.

sinkhorn_lr.LRSinkhornOutput(q, r, g, costs, ...)

Transport interface for a low-rank Sinkhorn solution.

Barycenter Solvers#

continuous_barycenter.FreeWassersteinBarycenter([...])

Continuous Wasserstein barycenter solver [Cuturi and Doucet, 2014].

continuous_barycenter.FreeBarycenterState([...])

Holds the state of the Wasserstein barycenter solver.

discrete_barycenter.FixedBarycenter([...])

A Wasserstein barycenter solver for histograms on a common geometry.

discrete_barycenter.SinkhornBarycenterOutput(f, ...)

Other Solvers#

univariate.UnivariateSolver([...])

Univariate solver to compute 1D OT distance over slices of data.

Sinkhorn Acceleration#

acceleration.Momentum([start, ...])

Momentum for Sinkhorn updates.

acceleration.AndersonAcceleration([memory, ...])

Implements Anderson acceleration for Sinkhorn.

Implicit Differentiation#

implicit_differentiation.ImplicitDiff([...])

Implicit differentiation of Sinkhorn algorithm.

implicit_differentiation.solve_jax_cg(lin, b)

Wrapper around JAX native linear solvers.

lineax_implicit.solve_lineax(lin, b[, ...])

Wrapper around lineax solvers.

Low-rank Sinkhorn Utilities#

lr_utils.unbalanced_dykstra_lse(c_q, c_r, ...)

Dykstra's algorithm for the unbalanced LRSinkhorn in LSE mode.

lr_utils.unbalanced_dykstra_kernel(k_q, k_r, ...)

Dykstra's algorithm for the unbalanced LRSinkhorn in kernel mode.