ott.geometry.pointcloud.PointCloud.vec_apply_cost#
- PointCloud.vec_apply_cost(arr, axis=0, fn=None)[source]#
Apply the geometry’s cost matrix in a vectorized way.
This function can be used when the cost matrix is squared euclidean and
fn
is a linear function.- Parameters:
- Return type:
- Returns:
A jnp.ndarray, [num_b, p] if axis=0 or [num_a, p] if axis=1