ott.geometry.low_rank.LRCGeometry.update_potential#

LRCGeometry.update_potential(f, g, log_marginal, iteration=None, axis=0)#

Carry out one Sinkhorn update for potentials, i.e. in log space.

Parameters
  • f (ndarray) – jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix

  • g (ndarray) – jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix

  • log_marginal (ndarray) – targeted marginal

  • iteration (Optional[int]) – used to compute epsilon from schedule, if provided.

  • axis (int) – axis along which the update should be carried out.

Return type

ndarray

Returns

new potential value, g if axis=0, f if axis is 1.