ott.geometry.geometry.Geometry.update_scaling

Geometry.update_scaling(scaling, marginal, iteration=None, axis=0)

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.