[MRG] Fix device placement in expected sliced plans

RELEASES.md

@@ -8,6 +8,7 @@
 
 #### Closed issues
 
+- Preserve input dtype and device for expected sliced plans, avoid materializing dense distance matrices for sparse plans, and fix weighted sparse-distance ordering (PR #846, Issue #845)
 - Fix the sign issue in updates of the previous transport plan in `ot.batch.proximal_bregman_log_plan_batch` (Issue #842)
 - Load triton before TensorFlow in `ot.backend` so that building a torch optimizer no longer segfaults the interpreter, and remove the `torch<2.12` pin from the doctest and documentation requirements (PR #839, Issue #816)
 - Fix mean centering in `ot.dr.fda` and `ot.dr.wda`: `np.mean(X)` returned a scalar instead of the per-feature mean, so `proj` did not center the data as documented. In `ot.dr.fda` the same pattern in the class means made the between-class scatter matrix independent of which features separate the classes, and FDA returned a non-discriminant direction (PR #840)
@@ -909,4 +910,4 @@ It provides the following solvers:
 * Optimal transport for domain adaptation with group lasso regularization
 * Conditional gradient and Generalized conditional gradient for regularized OT.
 
-Some demonstrations (both in Python and Jupyter Notebook format) are available in the examples folder.
+Some demonstrations (both in Python and Jupyter Notebook format) are available in the examples folder.

ot/sliced/_sliced_plans.py

@@ -12,7 +12,7 @@
 import warnings
 
 from ..backend import get_backend
-from ..utils import list_to_array, sparse_ot_dist, dist
+from ..utils import list_to_array, sparse_ot_dist
 from ._utils import get_random_projections
 from ..lp import wasserstein_1d
 from collections import namedtuple
@@ -103,9 +103,9 @@ def sliced_plans(
     m = X_t.shape[0]
 
     if a is None:
-        a = nx.ones(n) / n
+        a = nx.ones(n, type_as=X_s) / n
     if b is None:
-        b = nx.ones(m) / m
+        b = nx.ones(m, type_as=X_t) / m
 
     is_perm = (n == m) and (a == a.sum() / n).all() and (b == a.sum() / n).all()
 
@@ -493,7 +493,7 @@ def expected_sliced_plan(
     else:  # uniform weights
         if n_projections is None:
             n_projections = projections.shape[1]
-        weights = nx.ones(n_projections) / n_projections
+        weights = nx.ones(n_projections, type_as=X_s) / n_projections
 
     weights_e = nx.concatenate([plans[i].data * weights[i] for i in range(len(plans))])
     Xs_idx = nx.concatenate([plans[i].rows for i in range(len(plans))])
@@ -505,10 +505,16 @@ def expected_sliced_plan(
         plan = nx.todense(plan)
 
     if beta == 0.0:
-        if dense:
-            cost = nx.sum(plan * dist(X_s, X_t, metric=metric, p=p))
-        else:
-            cost = plan.multiply(dist(X_s, X_t, metric=metric, p=p)).sum()
+        cost = sparse_ot_dist(
+            X_s,
+            X_t,
+            Xs_idx,
+            Xt_idx,
+            weights_e,
+            metric=metric,
+            p=p,
+            batch_size=batch_size,
+        )
     if log:
         log_dict = {
             "projections": log_dict_plans["projections"],

ot/utils.py

@@ -391,7 +391,7 @@ def dist_idxs(idx_x1, idx_x2):
     for b in range(0, len(i), batch_size):
         d += nx.sum(
             dist_idxs(i[b : b + batch_size], j[b : b + batch_size])
-            * w[i[b : b + batch_size]]
+            * w[b : b + batch_size]
         )
 
     return d

test/sliced/test_sliced_plans.py

@@ -223,6 +223,40 @@ def test_expected_sliced():
     ot.sliced.expected_sliced_plan(x, y, a, b, projections=projections, batch_size=2)
 
 
+@pytest.mark.skipif(not torch, reason="PyTorch not installed")
+@pytest.mark.parametrize("dense", [True, False])
+@pytest.mark.parametrize("use_weights", [False, True])
+def test_expected_sliced_torch_preserves_device(dense, use_weights):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float64
+    x = torch.tensor(
+        [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]],
+        dtype=dtype,
+        device=device,
+    )
+    y = torch.tensor(
+        [[0.5, 0.0], [1.5, 0.0], [0.5, 1.0], [1.5, 1.0]],
+        dtype=dtype,
+        device=device,
+    )
+    projections = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=dtype, device=device)
+    if use_weights:
+        a = torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=dtype, device=device)
+        b = torch.tensor([0.4, 0.3, 0.2, 0.1], dtype=dtype, device=device)
+    else:
+        a = b = None
+
+    plan, cost = ot.sliced.expected_sliced_plan(
+        x, y, a, b, projections=projections, dense=dense
+    )
+
+    assert plan.device == x.device
+    assert cost.device == x.device
+    dense_plan = plan if dense else plan.to_dense()
+    expected_cost = torch.sum(dense_plan * ot.dist(x, y))
+    torch.testing.assert_close(cost, expected_cost)
+
+
 def test_sliced_plans_backends(nx):
     n = 10
     m = 24

test/test_utils.py

@@ -281,6 +281,20 @@ def test_dist():
         ot.dist(x, x, metric="fakeone")
 
 
+def test_sparse_ot_dist_uses_pair_weights():
+    x1 = np.array([[0.0], [1.0]])
+    x2 = np.array([[0.0], [5.0]])
+    rows = np.array([1, 0])
+    cols = np.array([0, 1])
+    weights = np.array([0.25, 0.75])
+
+    cost = ot.utils.sparse_ot_dist(
+        x1, x2, rows, cols, weights, metric="sqeuclidean", batch_size=1
+    )
+
+    np.testing.assert_allclose(cost, 0.25 * 1.0 + 0.75 * 25.0)
+
+
 @pytest.mark.parametrize("metric", lst_metrics)
 def test_dist_backends(nx, metric):
     n = 100