ott.geometry.low_rank.LRCGeometry.apply_square_cost#

LRCGeometry.apply_square_cost(arr, axis=0)[source]#

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

Return type

Array

Parameters

  • arr (jax._src.basearray.Array) –

  • axis (int) –