ott.geometry.low_rank.LRCGeometry.apply_square_cost

LRCGeometry.apply_square_cost(arr, axis=0)

Apply elementwise-square of cost matrix to array (vector or matrix).

Return type

ndarray

Parameters

• arr (jax._src.numpy.ndarray.ndarray) –

• axis (int) –