ott.geometry.grid.Grid.update_potential

ott.geometry.grid.Grid.update_potential#

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

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

Parameters:

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

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

  • log_marginal (Array) – 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:

Array

Returns:

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

Contents