`ot.sliced.expected_sliced_plan` fails on pytorch CUDA tensors

[tvercaut]

Describe the bug

ot.sliced.expected_sliced_plan() fails when inputs are PyTorch tensors on a CUDA device. The function appears to compute the transport plan on CUDA, but subsequently multiplies it with a distance matrix created on CPU, resulting in a device mismatch error.

To Reproduce

!pip install torch POT

import torch, ot
import sys, traceback

print(f"PyTorch Version: {torch.__version__}")
print(f"CUDA Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU Device: {torch.cuda.get_device_name(0)}")
print(f"POT Version: {ot.__version__}")

device = torch.device("cuda")
#device = torch.device("cpu")
X = torch.randn(4, 2, device=device)
w = torch.rand(4, device=device)
w = w / w.sum()
try:
  plan, cost = ot.sliced.expected_sliced_plan(X, X, n_projections=1)
except RuntimeError as e:
  print(f"\nexpected_sliced_plan with no weights failed:")
  traceback.print_exc(file=sys.stdout)
try:
  plan, cost = ot.sliced.expected_sliced_plan(X, X, w, w, n_projections=1, dense=False)
except RuntimeError as e:
  print(f"\nexpected_sliced_plan with dense=False failed:")
  traceback.print_exc(file=sys.stdout)
try:
  plan, cost = ot.sliced.expected_sliced_plan(X, X, w, w, n_projections=1, dense=True)
except RuntimeError as e:
  print(f"\nexpected_sliced_plan with dense=True failed:")
  traceback.print_exc(file=sys.stdout)

leads to

PyTorch Version: 2.11.0+cu128
CUDA Available: True
GPU Device: Tesla T4
POT Version: 0.9.7.post1

expected_sliced_plan with no weights failed:
Traceback (most recent call last):
  File "/tmp/ipykernel_757/207733201.py", line 18, in <cell line: 0>
    plan, cost = ot.sliced.expected_sliced_plan(X, X, n_projections=1)
                 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.13/dist-packages/ot/sliced/_sliced_plans.py", line 509, in expected_sliced_plan
    cost = nx.sum(plan * dist(X_s, X_t, metric=metric, p=p))
                  ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

expected_sliced_plan with dense=False failed:
Traceback (most recent call last):
  File "/tmp/ipykernel_757/207733201.py", line 23, in <cell line: 0>
    plan, cost = ot.sliced.expected_sliced_plan(X, X, w, w, n_projections=1, dense=False)
                 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.13/dist-packages/ot/sliced/_sliced_plans.py", line 511, in expected_sliced_plan
    cost = plan.multiply(dist(X_s, X_t, metric=metric, p=p)).sum()
           ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: mul: expected 'out' to be CUDA, but got CPU

expected_sliced_plan with dense=True failed:
Traceback (most recent call last):
  File "/tmp/ipykernel_757/207733201.py", line 28, in <cell line: 0>
    plan, cost = ot.sliced.expected_sliced_plan(X, X, w, w, n_projections=1, dense=True)
                 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.13/dist-packages/ot/sliced/_sliced_plans.py", line 509, in expected_sliced_plan
    cost = nx.sum(plan * dist(X_s, X_t, metric=metric, p=p))
                  ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Additional context

The issue occurs for the different cases I tried:

• implicit uniform weights,
• explicit weights with dense=True,
• explicit weights with dense=False.

The transport plan itself appears to be generated successfully. The failure apparently happens during the final cost computation due to mixed CPU/CUDA tensors.

[tvercaut]

Partially related to this issue, in the sparse return case, it may be preferable to avoid materialising a dense distance matrix to compute the cost:

POT/ot/sliced/_sliced_plans.py

Lines 508 to 511 in 9439b8a

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

I assume one could use ot.utils.sparse_ot_dist instead

[cananoo]

I plan to fix the device placement in expected_sliced_plan by making default weights and the assembled plan inherit the input tensor's dtype/device, and by using sparse_ot_dist for sparse cost evaluation instead of materializing a full distance matrix. I will add PyTorch regression coverage for dense and sparse plans with both implicit and explicit weights, with CUDA cases skipped when CUDA is unavailable, then run the sliced-plan tests and the repository lint/pre-commit checks.