[MRG] Preserve weighted low-rank plan shapes under integer indexing

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 (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)

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)
 

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(