From 84592b83e19838cc79fc4e1faec1ab8a9cf619da Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:31:04 -0700 Subject: [PATCH 1/3] Shift before exponentiating in coordinate_grad_semi_dual (closes #264) The semi-dual gradient computed exp(-(M[i] - beta) / reg) directly. Once beta grows, that overflows, the normalisation divides inf by inf, and khi comes back all-NaN, which is what makes the stochastic solvers unusable at the sizes reported in the issue. c_transform_entropic in the same module already guards against this by subtracting min(r) before exponentiating. This applies the same shift to coordinate_grad_semi_dual. It cancels in the ratio, so the returned gradient is unchanged wherever the old code did not overflow; verified to 1.7e-16 across a range of beta. Adds a test that the gradient is finite at large beta without even warning, and a test pinning the values against the unshifted formula in the regime where that formula is valid. --- RELEASES.md | 6 ++++++ ot/stochastic.py | 5 ++++- test/test_stochastic.py | 45 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index 063701229..755dde998 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -2,6 +2,12 @@ ## 0.9.8dev +#### Closed issues + +- Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #860, Issue #264) + +## 0.9.8dev + #### New features - Use `ot.utils.check_marginal` (and shape-tuple support in `ot.utils.unif`) to fill and validate default marginals consistently across solvers (Gromov, low-rank, stochastic, barycenter, factored) (PR #856) diff --git a/ot/stochastic.py b/ot/stochastic.py index 6df3ad771..8cbfe8f5e 100644 --- a/ot/stochastic.py +++ b/ot/stochastic.py @@ -67,7 +67,10 @@ def coordinate_grad_semi_dual(b, M, reg, beta, i): Advances in Neural Information Processing Systems (2016). """ r = M[i, :] - beta - exp_beta = np.exp(-r / reg) * b + # shift before exponentiating: the factor cancels in the ratio below, but + # without it exp overflows once beta grows and khi becomes all-NaN + min_r = np.min(r) + exp_beta = np.exp(-(r - min_r) / reg) * b khi = exp_beta / (np.sum(exp_beta)) return b - khi diff --git a/test/test_stochastic.py b/test/test_stochastic.py index 911d52b92..85e5a7701 100644 --- a/test/test_stochastic.py +++ b/test/test_stochastic.py @@ -13,6 +13,8 @@ # # License: MIT License +import warnings + import numpy as np import ot @@ -87,6 +89,49 @@ def test_stochastic_asgd(): # regularization term, a learning rate and a number of iteration +def test_coordinate_grad_semi_dual_large_beta(): + """Non-regression test for issue #264. + + The semi-dual gradient exponentiates -(M[i] - beta) / reg. Once beta grows + that overflows and the normalised khi becomes all-NaN, so the shift by the + minimum is what keeps the solver usable at large scale. + """ + rng = np.random.RandomState(0) + n_source, n_target, reg = 500, 60, 1 + M = ot.dist(rng.randn(n_source, 2), rng.randn(n_target, 2)) + b = ot.utils.unif(n_target) + + with warnings.catch_warnings(): + # narrow: unrelated warnings elsewhere must not make this brittle + warnings.simplefilter("error", RuntimeWarning) + grad = ot.stochastic.coordinate_grad_semi_dual( + b, M, reg, np.full(n_target, 800.0), 0 + ) + + assert not np.any(np.isnan(grad)) + # khi lives on the simplex, so the gradient b - khi sums to zero + np.testing.assert_allclose(np.sum(grad), 0.0, atol=1e-12) + + +def test_coordinate_grad_semi_dual_unchanged_without_overflow(): + """The shift cancels in the ratio, so benign inputs must be unaffected.""" + rng = np.random.RandomState(0) + n_source, n_target, reg = 200, 40, 1 + M = ot.dist(rng.randn(n_source, 2), rng.randn(n_target, 2)) + b = ot.utils.unif(n_target) + + for scale in (0.0, 0.5, 2.0): + beta = rng.randn(n_target) * scale + r = M[0, :] - beta + expected = b - (np.exp(-r / reg) * b) / np.sum(np.exp(-r / reg) * b) + np.testing.assert_allclose( + ot.stochastic.coordinate_grad_semi_dual(b, M, reg, beta, 0), + expected, + rtol=1e-12, + atol=1e-14, + ) + + def test_sag_asgd_sinkhorn(): # test all algorithms n = 10 From 016b47332f7160836a12774b29522dae2a6e6dad Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:51:58 -0700 Subject: [PATCH 2/3] Correct PR number in RELEASES.md --- RELEASES.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index 755dde998..81020e15a 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -4,7 +4,7 @@ #### Closed issues -- Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #860, Issue #264) +- Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264) ## 0.9.8dev @@ -21,7 +21,7 @@ #### Closed issues -- Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860) +- Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #866) - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859) - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854) - 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) From 0e07f1373c23813a12a5fcbd880b41f820b943e1 Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:52:27 -0700 Subject: [PATCH 3/3] Fold RELEASES entry into the existing 0.9.8dev section --- RELEASES.md | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index 81020e15a..96204df9c 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -2,12 +2,6 @@ ## 0.9.8dev -#### Closed issues - -- Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264) - -## 0.9.8dev - #### New features - Use `ot.utils.check_marginal` (and shape-tuple support in `ot.utils.unif`) to fill and validate default marginals consistently across solvers (Gromov, low-rank, stochastic, barycenter, factored) (PR #856) @@ -21,6 +15,7 @@ #### Closed issues +- Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264) - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #866) - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859) - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854)