From 0e8dfbcaae95ab3e3f1c36bb7d912f9d79478bcd Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Fri, 28 Aug 2026 12:11:19 +0200 Subject: [PATCH 1/2] sign correction --- ot/batch/_utils.py | 6 +++--- test/batch/test_solve_batch.py | 21 ++++++++++++++++----- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index 8bc96c3b9..d7b028f13 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -331,7 +331,7 @@ def proximal_bregman_log_plan_batch( .. math:: \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + (\textit{reg} + \textit{inner\_reg}) \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} - Denoting :math:`\mathbf{K}^{(k)} = - (\mathbf{C} + \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)})/(\textit{reg} + \textit{inner\_reg})`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: + Denoting :math:`\mathbf{K}^{(k)} = - (\mathbf{C} - \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)})/(\textit{reg} + \textit{inner\_reg})`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: .. math:: \mathbf{u}^{(i+1)} = \log(\mathbf{a}) - \text{LSE}(\mathbf{K}^{(k)} + \mathbf{v}^{(i)}) @@ -418,7 +418,7 @@ def proximal_bregman_log_plan_batch( log_T = nx.zeros(C.shape, type_as=C) for n_iters in range(max_iter): - K_proj = -(C + inner_reg * log_T) / (reg + inner_reg) + K_proj = -(C - inner_reg * log_T) / (reg + inner_reg) for _ in range(inner_iter): u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) @@ -433,7 +433,7 @@ def proximal_bregman_log_plan_batch( break if grad == "last_step": - K_proj = -(C_ + inner_reg * log_T) / (reg + inner_reg) + K_proj = -(C_ - inner_reg * log_T) / (reg + inner_reg) for _ in range(inner_iter): u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index 08e5ff81e..ee2d35c51 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -51,6 +51,7 @@ def test_solve_batch_vs_solve(reg, method, reg_type): base_plan[i] = res_i.plan base_value[i] = res_i.value_linear + inner_reg = 1e-1 if reg is None or reg == 0 else 1e-3 res = solve_batch( C, max_iter=10000, @@ -59,11 +60,15 @@ def test_solve_batch_vs_solve(reg, method, reg_type): reg=reg, method=method, reg_type=reg_type, - inner_reg=1e-3, + inner_reg=inner_reg, ) plan = res.plan value = res.value_linear - np.testing.assert_allclose(plan, base_plan, atol=tol * 10) + if reg is None or reg == 0: + np.testing.assert_allclose(plan.sum(axis=2), 1 / n, atol=tol * 10) + np.testing.assert_allclose(plan.sum(axis=1), 1 / d, atol=tol * 10) + else: + np.testing.assert_allclose(plan, base_plan, atol=tol * 10) np.testing.assert_allclose(value, base_value, atol=tol * 10) @@ -76,16 +81,19 @@ def test_backend_proximal_bregman_log_plan_batch(nx, reg, inner_iter): d = 7 rng = np.random.RandomState(0) C = rng.rand(batchsize, n, d) + inner_reg = 1e-1 if reg is None or reg == 0 else 1e-3 res = proximal_bregman_log_plan_batch( nx.from_numpy(C), reg=reg, - inner_reg=1e-3, + inner_reg=inner_reg, max_iter=10000, tol=tol, inner_iter=inner_iter, grad="detach", ) plan = nx.to_numpy(res["T"]) + np.testing.assert_allclose(plan.sum(axis=2), 1 / n, atol=tol * 10) + np.testing.assert_allclose(plan.sum(axis=1), 1 / d, atol=tol * 10) for i in range(batchsize): C_i = C[i] res_i = solve( @@ -93,8 +101,11 @@ def test_backend_proximal_bregman_log_plan_batch(nx, reg, inner_iter): reg=reg, tol=tol, ) - plan_i = res_i.plan - np.testing.assert_allclose(plan_i, plan[i], atol=tol * 10) + if reg is None or reg == 0: + value = np.sum(C_i * plan[i]) + np.testing.assert_allclose(res_i.value_linear, value, atol=tol * 10) + else: + np.testing.assert_allclose(res_i.plan, plan[i], atol=tol * 10) def test_bregman_batch(): From cd8cd364eba537562b64b5e43543aabcecfdddf4 Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Fri, 28 Aug 2026 12:21:30 +0200 Subject: [PATCH 2/2] add issur reference in the release.md --- RELEASES.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index 9a5ddef50..526f6409e 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -8,6 +8,7 @@ #### Closed issues +- 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) - `ot.dr.fda` and `ot.dr.wda` no longer modify the input array `X` in place (PR #840) @@ -908,4 +909,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. \ No newline at end of file