ot_markov_distances.sinkhorn module

Differentiable sinkhorn divergence. This module provides a pytorch/autograd compatible implementation of sinkhorn divergence, using the method described in [FSejourneV+19]

ot_markov_distances.sinkhorn.sinkhorn_internal(a: Tensor, b: Tensor, C: Tensor, epsilon: float, k: int = 100, *, check_convergence_interval: int | float = 0.1, cv_atol=0.0001, cv_rtol=1e-05, return_has_converged: Literal[False] = False) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]
ot_markov_distances.sinkhorn.sinkhorn_internal(a: Tensor, b: Tensor, C: Tensor, epsilon: float, k: int = 100, *, check_convergence_interval: int | float = 0.1, cv_atol=0.0001, cv_rtol=1e-05, return_has_converged: Literal[True]) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]

Same as sinkhorn, but returns f, g and P instead of the result Beware, this function does not have the shortcut differentiation

(It can still be differentiated by autograd though)

Parameters:

  • a – (*batch, n) vector of the first distribution

  • b – (*batch, m) vector of the second distribtion

  • C – (*batch, n, m) cost matrix

  • epsilon – regularisation term for sinkhorn

  • k – max number of sinkhorn iterations (default 100)

  • check_convergence_interval – if int, check for convergence every check_convergence_interval. If float, check for convergence every check_convergence_interval * k. If 0, never check for convergence (apart from the last iteration if return_has_converged==True) If convergence is reached early, the algorithm returns.

  • cv_atol – absolute and relative tolerance for the converegence criterion

  • cv_rtol – absolute and relative tolerance for the converegence criterion

  • return_has_converged – whether to return a boolean indicating whether the algorithm has converged. Setting this to True means that the function will always check for convergence at the last iteration (regardless of the value of check_convergence_interval)

Returns:

  • f – Tensor (*batch, n)

  • g – Tensor (*batch, m)

  • log_P – Tensor (*batch, n, m)

  • has_converged – bool only returened if return_has_converged is True. Indicates whether the algorithm has converged

ot_markov_distances.sinkhorn.sinkhorn(a: Tensor, b: Tensor, C: Tensor, epsilon: float, max_iter: int = 100, *, check_convergence_interval: int | float = 0.1, cv_atol=0.0001, cv_rtol=1e-05, return_has_converged: Literal[True, False] = False) → Tensor

Differentiable sinkhorn distance

This is a pytorch implementation of sinkhorn, batched (over a, b and C)

It is compatible with pytorch autograd gradient computations.

See the documentation of Sinkhorn for details.

Parameters:

  • a – (*batch, n) vector of the first distribution

  • b – (*batch, m) vector of the second distribtion

  • C – (*batch, n, m) cost matrix

  • epsilon – regularisation term for sinkhorn

  • max_iter – max number of sinkhorn iterations (default 100)

  • check_convergence_interval – if int, check for convergence every check_convergence_interval. If float, check for convergence every check_convergence_interval * max_iter. If 0, never check for convergence (apart from the last iteration if return_has_converged==True) If convergence is reached early, the algorithm returns.

  • cv_atol – absolute and relative tolerance for the converegence criterion

  • cv_rtol – absolute and relative tolerance for the converegence criterion

  • return_has_converged – whether to return a boolean indicating whether the algorithm has converged. Setting this to True means that the function will always check for convergence at the last iteration (regardless of the value of check_convergence_interval)

Returns:

Tensor – (*batch). result of the sinkhorn computation