From 6651f1cb30ede239912a5b8239b7b1d4dd1b293c Mon Sep 17 00:00:00 2001 From: tommoral Date: Wed, 9 Sep 2026 19:40:50 +0000 Subject: [PATCH 1/3] FIX device placement in batch Sinkhorn so solve_batch runs on GPU --- RELEASES.md | 1 + ot/batch/_utils.py | 16 ++++++++-------- test/batch/test_solve_batch.py | 19 +++++++++++++++++++ 3 files changed, 28 insertions(+), 8 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index 6ead2ce34..932b4a679 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -8,6 +8,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 #847) - 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) diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index d7b028f13..f8e451ae7 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -137,17 +137,17 @@ def bregman_projection_batch( B, n, m = K.shape if a is None: - a = nx.ones((B, n)) / n + a = nx.ones((B, n), type_as=K) / n if b is None: - b = nx.ones((B, m)) / m + b = nx.ones((B, m), type_as=K) / m 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) @@ -261,9 +261,9 @@ def bregman_log_projection_batch( B, n, m = K.shape if a is None: - a = nx.ones((B, n)) / n + a = nx.ones((B, n), type_as=K) / n if b is None: - b = nx.ones((B, m)) / m + b = nx.ones((B, m), type_as=K) / m 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() @@ -398,9 +398,9 @@ def proximal_bregman_log_plan_batch( B, n, m = C.shape if a is None: - a = nx.ones((B, n)) / n + a = nx.ones((B, n), type_as=C) / n if b is None: - b = nx.ones((B, m)) / m + b = nx.ones((B, m), type_as=C) / m if reg is None: reg = 0.0 diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index ee2d35c51..ff2288e41 100644 --- a/test/batch/test_solve_batch.py +++ b/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 From cabcb11defc7398b99449aacdb3d21d3685ded25 Mon Sep 17 00:00:00 2001 From: tommoral Date: Wed, 9 Sep 2026 20:01:18 +0000 Subject: [PATCH 2/3] ENH add ot.utils.check_marginal and unif shape support, use in batch solvers --- ot/batch/_linear.py | 8 +++--- ot/batch/_quadratic.py | 8 +++--- ot/batch/_utils.py | 19 +++++-------- ot/utils.py | 60 ++++++++++++++++++++++++++++++++++++------ test/test_utils.py | 26 ++++++++++++++++++ 5 files changed, 91 insertions(+), 30 deletions(-) diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index 3c07ae639..59c3e5c5a 100644 --- a/ot/batch/_linear.py +++ b/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 diff --git a/ot/batch/_quadratic.py b/ot/batch/_quadratic.py index ccb8a8792..ef512e298 100644 --- a/ot/batch/_quadratic.py +++ b/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( diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index f8e451ae7..1521e663f 100644 --- a/ot/batch/_utils.py +++ b/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,10 +137,8 @@ def bregman_projection_batch( B, n, m = K.shape - if a is None: - a = nx.ones((B, n), type_as=K) / n - if b is None: - b = nx.ones((B, m), type_as=K) / 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) @@ -260,10 +259,8 @@ def bregman_log_projection_batch( B, n, m = K.shape - if a is None: - a = nx.ones((B, n), type_as=K) / n - if b is None: - b = nx.ones((B, m), type_as=K) / 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), type_as=C) / n - if b is None: - b = nx.ones((B, m), type_as=C) / 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 diff --git a/ot/utils.py b/ot/utils.py index 5fee6109c..6336db07a 100644 --- a/ot/utils.py +++ b/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): diff --git a/test/test_utils.py b/test/test_utils.py index dce79af6a..23f1ca96a 100644 --- a/test/test_utils.py +++ b/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 From 57ec105a2486da39b643211353e2a4faabffa767 Mon Sep 17 00:00:00 2001 From: tommoral Date: Wed, 9 Sep 2026 20:07:42 +0000 Subject: [PATCH 3/3] DOC fix PR number in RELEASES entry --- RELEASES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index 932b4a679..ac6206c70 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -8,7 +8,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 #847) +- 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) - 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)