Skip to content

Fix Grid.to_LRCGeometry for grids with 2 points along a dimension - #711

Open
raashish1601 wants to merge 1 commit into
ott-jax:mainfrom
raashish1601:fix/669-grid-lrc-side-two
Open

raashish1601 wants to merge 1 commit into
ott-jax:mainfrom
raashish1601:fix/669-grid-lrc-side-two

Conversation

@raashish1601

Copy link
Copy Markdown

Fixes #669.

Grid.to_LRCGeometry() converts the point cloud of each grid dimension with geom.to_LRCGeometry(rank=0, ...) and then reads cost_1 / cost_2. For a squared Euclidean (or negative dot product) point cloud, PointCloud.to_LRCGeometry() returns the point cloud itself when the factors would not be smaller than the cost matrix (n * m <= (n + m) * d). With 2 points on a dimension that is always the case (4 <= 4), so any grid with a dimension of size 2 fails with AttributeError: 'PointCloud' object has no attribute 'cost_1'. SinkhornOutput.primal_cost hits this through transport_cost_at_geom(), since Grid.can_LRC is True.

When the slice conversion does not return an LRCGeometry, the grid now factorizes that slice with the exact SVD from Geometry.to_LRCGeometry(rank=0), which is what the comment above the loop already asks for.

Tests:

  • test_conversion_grid in tests/geometry/lr_cost_test.py is parametrized with grids that have size-2 dimensions ([2], [2, 3], [3, 2, 2]) in addition to the existing [6, 7, 11].
  • test_primal_cost_grid_with_two_points in tests/solvers/linear/sinkhorn_grid_test.py compares primal_cost on grids of size (2,), (2, 2) and (2, 3) with the same problem on a PointCloud.

Both fail on main and pass with the change. tests/solvers/linear/sinkhorn_grid_test.py and the grid conversion tests pass locally (CPU), and ruff, isort and yapf at the pre-commit versions report no changes.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

SinkhornOutput.primal_cost crashes for grid of side length 2

1 participant