Fix/low rank conversion pointcloud

[clijo]

This PR fixes the issue #669 .

root cause

Root cause seems to be that PointCloud._check_LRC_dim was False in this case, and so PointCloud.to_LRCGeometry() returned the unmodified point cloud; which meant that the geom.to_LRCGeometry(rank=0, scale=scale, **kwargs) call in grid.py was not returning a low_rank.LRCGeometry object, and so the line c_1, c_2 = geom.cost_1, geom.cost_2 causes the crash because AttributeError: 'PointCloud' object has no attribute 'cost_1'.

fix

As the inline comments above geom.to_LRCGeometry(rank=0, scale=scale, **kwargs) call in grid.py explain, rank=0 is supposed to force low rank conversion. But, PointCloud.to_LRCGeometry() was ignoring that kwarg. I added a Boolean flag called force = (kwargs.get("rank") == 0) to coerce a low rank conversion.

The repro code from the issue now runs:

import jax.numpy as jnp
from ott.geometry.grid import Grid
from ott.problems.linear.linear_problem import LinearProblem
from ott.solvers.linear.sinkhorn import Sinkhorn
grid = Grid(grid_size=(2,))
prob = LinearProblem(grid, a=jnp.array([1,0]), b=jnp.array([0,1]))
out = Sinkhorn()(prob)
print(out.primal_cost)

1.0

[codecov]

Codecov Report

❌ Patch coverage is 66.66667% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 87.32%. Comparing base (2f93834) to head (6a0f4dd).

Files with missing lines	Patch %	Lines
src/ott/geometry/pointcloud.py	66.66%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #672   +/-   ##
=======================================
  Coverage   87.32%   87.32%           
=======================================
  Files          82       82           
  Lines        8480     8481    +1     
  Branches      581      581           
=======================================
+ Hits         7405     7406    +1     
  Misses        924      924           
  Partials      151      151           

Files with missing lines	Coverage Δ
src/ott/geometry/pointcloud.py	94.09% <66.66%> (+0.02%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.