ott.solvers.linear.sinkhorn_lr.LRSinkhorn.one_iteration

ott.solvers.linear.sinkhorn_lr.LRSinkhorn.one_iteration#

LRSinkhorn.one_iteration(ot_prob, state, iteration, compute_error)[source]#

Carries out one low-rank Sinkhorn iteration.

Depending on lse_mode, these iterations can be either in:

  • log-space for numerical stability.

  • scaling space, using standard kernel-vector multiply operations.

Parameters:
  • ot_prob (LinearProblem) – the transport problem definition

  • state (LRSinkhornState) – LRSinkhornState named tuple.

  • iteration (int) – the current iteration of the Sinkhorn outer loop.

  • compute_error (bool) – flag to indicate this iteration computes/stores an error

Return type:

LRSinkhornState

Returns:

The updated state.