OTT tools: A set of tools to use OT in differentiable ML pipelines.

The tools package contains high level functions that build on outputs produced by core functions. They can be used to compute Sinkhorn divergences 1, instantiate transport matrices, provide differentiable approximations to ranks and quantile functions 2, etc.

Optimal Transport#

transport.Transport([problem, solver_output])

Implements a core.problems.Transport interface to transport solutions.

Sinkhorn Divergence#


Computes Sinkhorn divergence defined by a geometry, weights, parameters.

Soft Sorting algorithms#

soft_sort.quantile(inputs[, axis, level, weight])

Applies the soft quantile operator on the input tensor.

soft_sort.quantile_normalization(inputs, targets)

Renormalizes inputs so that its quantiles match those of targets/weights.

soft_sort.quantize(inputs[, num_levels, axis])

Soft quantizes an input according using num_levels values along axis.

soft_sort.ranks(inputs[, axis, num_targets])

Applies the soft trank operator on input tensor.

soft_sort.sort(inputs[, axis, topk, num_targets])

Applies the soft sort operator on a given axis of the input.

soft_sort.sort_with(inputs, criterion[, topk])

Sort a multidimensional array according to a real valued criterion.


  1. Séjourné et al., Sinkhorn Divergences for Unbalanced Optimal Transport, arXiv:1910.12958.

  1. Cuturi et al. Differentiable Ranking and Sorting using Optimal Transport, NeurIPS’19.