From 125061e6d7d6e2cc1f07dbc5e9e570596d94bf83 Mon Sep 17 00:00:00 2001 From: ktz03 <2484593937@qq.com> Date: Mon, 19 Oct 2026 11:42:34 +0800 Subject: [PATCH 1/2] Fix weighted low-rank lazy tensor integer indexing --- RELEASES.md | 1 + ot/utils.py | 2 +- test/test_utils.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index 8b5aff51d..62a714afc 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -18,6 +18,7 @@ #### Closed issues +- Fix integer row and scalar indexing of weighted low-rank lazy transport plans returning an extra leading dimension (Issue #892). - Fix internal arrays created on the torch default device instead of the input device in `ot.partial.partial_wasserstein`, `ot.wasserstein_circle` (and `ot.sliced_wasserstein_sphere`), the high-dimensional Bures-Wasserstein functions in `ot.gaussian`, `ot.gmm`, `ot.lowrank_sinkhorn` with `init="deterministic"`, the semi-relaxed (F)GW barycenters and `ot.utils.projection_sparse_simplex`, which made them fail on GPU inputs. Add a `meta_default_device` test fixture to catch such allocations on CPU (PR #883, Issue #852) - Allow `NumpyBackend.seed` to adopt an existing `np.random.RandomState` instance and remove NumPy-specific random sampling paths in sliced utilities (PR #849, Issue #848) - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860) diff --git a/ot/utils.py b/ot/utils.py index 69a4854b5..6fbd6973d 100644 --- a/ot/utils.py +++ b/ot/utils.py @@ -898,7 +898,7 @@ def func(i, j, Q, R): else: def func(i, j, Q, R, d): - return nx.dot(Q[i] * d[None, :], R[j].T) + return nx.dot(Q[i] * d, R[j].T) T = LazyTensor(shape, func, Q=Q, R=R, d=d) diff --git a/test/test_utils.py b/test/test_utils.py index 9c93d520a..92850382b 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -739,6 +739,35 @@ def test_lowrank_LazyTensor(nx): np.testing.assert_allclose(nx.to_numpy(T[:]), nx.to_numpy(T0)) +@pytest.mark.parametrize("weighted", [False, True]) +@pytest.mark.parametrize( + "index", + [ + 0, + -1, + (1, 2), + (-1, -2), + (1, slice(1, None)), + (slice(None), 1), + (slice(1, 3), slice(None)), + slice(None), + ], +) +def test_lowrank_LazyTensor_indexing(nx, weighted, index): + Q, R, d = nx.from_numpy( + np.arange(10, dtype=float).reshape(5, 2), + np.arange(6, dtype=float).reshape(3, 2), + np.array([0.25, 2.0]), + ) + if not weighted: + d = None + T = ot.utils.get_lowrank_lazytensor(Q, R, d) + dense = nx.dot(Q if d is None else Q * d, R.T) + + assert T[index].shape == dense[index].shape + np.testing.assert_allclose(nx.to_numpy(T[index]), nx.to_numpy(dense[index])) + + def test_labels_to_mask_helper(nx): y = np.array([1, 0, 2, 2, 1]) out = np.array( From d1bb591d33e2b7c6df34a0a3abe772350203505d Mon Sep 17 00:00:00 2001 From: ktz03 <2484593937@qq.com> Date: Mon, 19 Oct 2026 11:46:29 +0800 Subject: [PATCH 2/2] Link the low-rank indexing fix in release notes --- RELEASES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index 62a714afc..7d3b0bf0b 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -18,7 +18,7 @@ #### Closed issues -- Fix integer row and scalar indexing of weighted low-rank lazy transport plans returning an extra leading dimension (Issue #892). +- Fix integer row and scalar indexing of weighted low-rank lazy transport plans returning an extra leading dimension (PR #893, Issue #892). - Fix internal arrays created on the torch default device instead of the input device in `ot.partial.partial_wasserstein`, `ot.wasserstein_circle` (and `ot.sliced_wasserstein_sphere`), the high-dimensional Bures-Wasserstein functions in `ot.gaussian`, `ot.gmm`, `ot.lowrank_sinkhorn` with `init="deterministic"`, the semi-relaxed (F)GW barycenters and `ot.utils.projection_sparse_simplex`, which made them fail on GPU inputs. Add a `meta_default_device` test fixture to catch such allocations on CPU (PR #883, Issue #852) - Allow `NumpyBackend.seed` to adopt an existing `np.random.RandomState` instance and remove NumPy-specific random sampling paths in sliced utilities (PR #849, Issue #848) - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860)