[MRG] Fix device placement in batch Sinkhorn (solve_batch on GPU)

RELEASES.md

@@ -11,6 +11,7 @@
 
 #### Closed issues
 
+- Fix device placement in `ot.batch.bregman_projection_batch` so `ot.solve_batch(..., method="sinkhorn")` no longer crashes on GPU when the torch default device is CPU (PR #851)
 - 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)
 - Build the CUDA generator and the CUDA entries of `TorchBackend.__type_list__` lazily, so that using POT with CPU-only torch tensors no longer initialises a CUDA context and claims device memory (PR #847, Issue #612)
 - Fix the sign issue in updates of the previous transport plan in `ot.batch.proximal_bregman_log_plan_batch` (Issue #842)

ot/batch/_linear.py

@@ -10,7 +10,7 @@
 # License: MIT License
 
 from ..backend import get_backend
-from ..utils import OTResult
+from ..utils import OTResult, check_marginal
 from ._utils import (
     bregman_log_projection_batch,
     bregman_projection_batch,
@@ -375,10 +375,8 @@ def solve_batch(
 
     B, n, m = M.shape
 
-    if a is None:
-        a = nx.ones((B, n), type_as=M) / n
-    if b is None:
-        b = nx.ones((B, m), type_as=M) / m
+    a = check_marginal(a, (B, n), type_as=M, nx=nx)
+    b = check_marginal(b, (B, m), type_as=M, nx=nx)
 
     if method == "log_sinkhorn":
         K = -M / reg

ot/batch/_quadratic.py

@@ -12,7 +12,7 @@
 from ot.backend import get_backend
 from ot.batch._linear import loss_linear_batch
 from ot.batch._utils import bmv, bop, bregman_log_projection_batch
-from ot.utils import deprecated, list_to_array
+from ot.utils import check_marginal, deprecated, list_to_array
 
 
 def tensor_batch(
@@ -621,10 +621,8 @@ def solve_gromov_batch(
     nx = get_backend(a, b, M, Ca, Cb, T_init)
     B, n, m = (Ca.shape[0], Ca.shape[1], Cb.shape[1])
 
-    if a is None:
-        a = nx.ones((B, n), type_as=Ca) / n
-    if b is None:
-        b = nx.ones((B, m), type_as=Cb) / m
+    a = check_marginal(a, (B, n), type_as=Ca, nx=nx)
+    b = check_marginal(b, (B, m), type_as=Cb, nx=nx)
 
     if symmetric is None:
         symmetric = nx.allclose(Ca, transpose(Ca, nx=nx), atol=1e-10) and nx.allclose(

ot/batch/_utils.py

@@ -10,6 +10,7 @@
 # License: MIT License
 
 from ot.backend import get_backend
+from ot.utils import check_marginal
 
 
 def entropy_batch(T, nx=None, eps=1e-16):
@@ -136,18 +137,16 @@ def bregman_projection_batch(
 
     B, n, m = K.shape
 
-    if a is None:
-        a = nx.ones((B, n)) / n
-    if b is None:
-        b = nx.ones((B, m)) / m
+    a = check_marginal(a, (B, n), type_as=K, nx=nx)
+    b = check_marginal(b, (B, m), type_as=K, nx=nx)
 
     if grad == "detach":
         K = nx.detach(K)
     elif grad == "last_step":
         K_, K = K.clone(), nx.detach(K)
 
-    f = nx.ones((B, n))  # a / nx.sum(K, axis=2)
-    g = nx.ones((B, m))  # b / nx.sum(K, axis=1)
+    f = nx.ones((B, n), type_as=K)  # a / nx.sum(K, axis=2)
+    g = nx.ones((B, m), type_as=K)  # b / nx.sum(K, axis=1)
 
     for n_iters in range(max_iter):
         f = a / nx.sum(K * g[:, None, :], axis=2)
@@ -260,10 +259,8 @@ def bregman_log_projection_batch(
 
     B, n, m = K.shape
 
-    if a is None:
-        a = nx.ones((B, n)) / n
-    if b is None:
-        b = nx.ones((B, m)) / m
+    a = check_marginal(a, (B, n), type_as=K, nx=nx)
+    b = check_marginal(b, (B, m), type_as=K, nx=nx)
 
     u = nx.zeros((B, n), type_as=K)  # u = nx.log(a) - nx.logsumexp(K, axis=2).squeeze()
     v = nx.zeros((B, m), type_as=K)  # v = nx.log(b) - nx.logsumexp(K, axis=1).squeeze()
@@ -397,10 +394,8 @@ def proximal_bregman_log_plan_batch(
 
     B, n, m = C.shape
 
-    if a is None:
-        a = nx.ones((B, n)) / n
-    if b is None:
-        b = nx.ones((B, m)) / m
+    a = check_marginal(a, (B, n), type_as=C, nx=nx)
+    b = check_marginal(b, (B, m), type_as=C, nx=nx)
 
     if reg is None:
         reg = 0.0

ot/utils.py

@@ -236,27 +236,71 @@ def projection_sparse_simplex(V, max_nz, z=1, axis=None, nx=None):
         return projection_sparse_simplex(V, max_nz, z, axis=1).ravel()
 
 
-def unif(n, type_as=None):
+def unif(shape, type_as=None):
     r"""
-    Return a uniform histogram of length `n` (simplex).
+    Return a uniform histogram normalized over its last dimension (simplex).
 
     Parameters
     ----------
-    n : int
-        number of bins in the histogram
+    shape : int or tuple of int
+        Number of bins in the histogram, or the full output shape. An integer
+        ``n`` is equivalent to ``(n,)``. For a shape ``(..., n)`` the output is
+        normalized over the last dimension, e.g. ``(B, n)`` returns ``B``
+        uniform histograms of length ``n``.
     type_as : array-like
         array of the same type of the expected output (numpy/pytorch/jax)
 
     Returns
     -------
-    h : array-like, shape (n,)
-        histogram of length `n` such that :math:`\forall i, \mathbf{h}_i = \frac{1}{n}`
+    h : array-like, shape ``shape``
+        uniform histogram(s) such that each slice along the last dimension sums
+        to one, i.e. every entry equals :math:`\frac{1}{n}` with ``n`` the last
+        dimension.
     """
+    size = (int(shape),) if isinstance(shape, (int, np.integer)) else tuple(shape)
+    n = size[-1]
     if type_as is None:
-        return np.ones((n,)) / n
+        return np.ones(size) / n
     else:
         nx = get_backend(type_as)
-        return nx.ones((n,), type_as=type_as) / n
+        return nx.ones(size, type_as=type_as) / n
+
+
+def check_marginal(a, shape, type_as=None, nx=None):
+    r"""Validate or fill a marginal :math:`\mathbf{a}` against an expected shape.
+
+    When ``a`` is ``None`` it is filled with the uniform marginal of the given
+    ``shape`` on the backend, device and dtype of ``type_as`` (via :func:`unif`).
+    When provided, its shape is checked against ``shape`` and, if ``type_as`` is
+    given, it must share ``type_as``'s dtype and device -- otherwise an explicit
+    error is raised instead of a deep error later in the solver.
+
+    Parameters
+    ----------
+    a : array-like, shape ``shape``, or None
+        The marginal to validate, or None to fill with a uniform marginal.
+    shape : int or tuple of int
+        Expected shape of the marginal (see :func:`unif`).
+    type_as : array-like, optional
+        Array fixing the backend, device and dtype of the filled/validated marginal.
+    nx : backend object, optional
+        Numerical backend to use. If None, it is inferred from ``type_as``.
+
+    Returns
+    -------
+    a : array-like, shape ``shape``
+        The validated marginal.
+    """
+    size = (int(shape),) if isinstance(shape, (int, np.integer)) else tuple(shape)
+    if a is None:
+        return unif(size, type_as=type_as)
+    if tuple(a.shape) != size:
+        raise ValueError(f"marginal has shape {tuple(a.shape)}, expected {size}")
+    if type_as is not None:
+        if nx is None:
+            nx = get_backend(type_as)
+        nx.assert_same_dtype_device(type_as, a)
+    return a
 
 
 def clean_zeros(a, b, M):

test/batch/test_solve_batch.py

@@ -190,3 +190,22 @@ def test_backend(nx, method):
     M = dist_batch(X, X)
     solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, method=method)
     solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5, method=method)
+
+
+@pytest.mark.skipif(not torch, reason="torch not installed")
+@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"])
+def test_solve_batch_device(method):
+    """Solve on each available device without internal CPU/GPU mismatch."""
+    batchsize = 2
+    n = 4
+    d = 2
+    X = np.random.randn(batchsize, n, d)
+
+    devices = [torch.device("cpu")]
+    if torch.cuda.is_available():
+        devices.append(torch.device("cuda"))
+    for device in devices:
+        Xd = torch.tensor(X, device=device)
+        M = dist_batch(Xd, Xd)
+        res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, method=method)
+        assert res.plan.device == Xd.device

test/test_utils.py

@@ -204,6 +204,32 @@ def test_unif_backend(nx):
         np.testing.assert_allclose(1, np.sum(nx.to_numpy(u)), atol=1e-6)
 
 
+def test_unif_shape():
+    # int form is unchanged
+    np.testing.assert_allclose(1, np.sum(ot.unif(10)))
+    # shape-tuple form normalizes over the last dimension
+    u = ot.unif((3, 5))
+    assert u.shape == (3, 5)
+    np.testing.assert_allclose(1, np.sum(u, axis=-1))
+
+
+def test_check_marginal(nx):
+    from ot.utils import check_marginal
+
+    M = nx.from_numpy(np.random.rand(2, 4, 5))
+    # None -> uniform on the reference backend/device, normalized over last dim
+    a = check_marginal(None, (2, 4), type_as=M)
+    b = check_marginal(None, (2, 5), type_as=M)
+    assert tuple(a.shape) == (2, 4) and tuple(b.shape) == (2, 5)
+    np.testing.assert_allclose(1, nx.to_numpy(nx.sum(a, axis=-1)), atol=1e-6)
+    # a valid provided marginal is returned unchanged
+    a2 = nx.from_numpy(np.full((2, 4), 1 / 4))
+    assert check_marginal(a2, (2, 4), type_as=M) is a2
+    # wrong shape raises a clear error
+    with pytest.raises(ValueError):
+        check_marginal(a2, (2, 5), type_as=M)
+
+
 def test_dist():
     n = 10