ott.geometry.geometry.Geometry.update_scaling#

Geometry.update_scaling(scaling, marginal, iteration=None, axis=0)[source]#

Carry out one Sinkhorn update for scalings, using kernel directly.

Parameters

  • scaling (Array) – jnp.ndarray of num_a or num_b positive values.

  • 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 scaling vector, of size num_b if axis=0, num_a if axis is 1.