From daaaa2706be66223ca58ed3a5e4da8f3b1a40adf Mon Sep 17 00:00:00 2001 From: Jonathan Legrand Date: Mon, 21 Sep 2026 11:02:29 +0200 Subject: [PATCH 1/7] Add overflow challenge for T_rand --- test/test_gmm.py | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/test/test_gmm.py b/test/test_gmm.py index 629a68d57..5e5c928ca 100644 --- a/test/test_gmm.py +++ b/test/test_gmm.py @@ -6,6 +6,7 @@ # # License: MIT License +import warnings import numpy as np import pytest from ot.utils import proj_simplex @@ -160,6 +161,48 @@ def test_gmm_apply_map(): gmm_ot_apply_map(x, m_s, m_t, C_s, C_t, w_s, w_t, plan=plan) +def test_gmm_apply_map_overflow(nx): + # In the original implementation + # log_diff computes log(g_i) - log(g_j) = log(g_i/g_j) + # so that g_k / sum_i(g_i) is easier to compute afterwards + # Problem is that if g_j(x) is small and g_i(x) is bigger, + # we can end up with high log ratios which overflows the + # subsequent exp calculation and throws a very unpleasant warning + x_coord = 12 + d = 1 + k = 2 + x = np.array([x_coord]).reshape((-1, d)) + + m_s = np.array([x_coord, 0]).reshape((k, d)) + m_t = m_s.copy() + + C_s = np.ones(m_s.shape).reshape((k, d, d)) / 10 + C_t = C_s.copy() + + w_s = np.array([0.5, 0.5]) + w_t = w_s.copy() + m_s = nx.from_numpy(m_s) + m_t = nx.from_numpy(m_t) + C_s = nx.from_numpy(C_s) + C_t = nx.from_numpy(C_t) + w_s = nx.from_numpy(w_s) + w_t = nx.from_numpy(w_t) + + with warnings.catch_warnings(): + warnings.simplefilter(action="error") + gmm_ot_apply_map( + x, + m_s, + m_t, + C_s, + C_t, + w_s, + w_t, + method="rand", + seed=0, + ) + + @pytest.mark.skipif(not torch, reason="No torch available") def test_gradient_gmm_ot_loss_pytorch(): m_s, m_t, C_s, C_t, w_s, w_t = get_gmms() From 9172eb95d7d74d23b173408b635f0d9041a385b0 Mon Sep 17 00:00:00 2001 From: Jonathan Legrand Date: Mon, 21 Sep 2026 11:42:34 +0200 Subject: [PATCH 2/7] Implement logsumexp trick --- ot/gmm.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/ot/gmm.py b/ot/gmm.py index 3cb6cea6a..7c2c42e7c 100644 --- a/ot/gmm.py +++ b/ot/gmm.py @@ -249,6 +249,16 @@ def gmm_ot_plan(m_s, m_t, C_s, C_t, w_s, w_t, log=False): return emd(w_s, w_t, D, log=log) +def logsumexp(x, scaling_factor=None): + """ + logsumexp trick such as in https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/ + """ + nx = get_backend(x, scaling_factor) + shift = nx.max(x) + y = shift + nx.log(nx.sum(scaling_factor * nx.exp(x - shift))) + return y + + def gmm_ot_apply_map( x, m_s, m_t, C_s, C_t, w_s, w_t, plan=None, method="bary", seed=None ): @@ -354,10 +364,10 @@ def gmm_ot_apply_map( for i_sample in range(n_samples): log_g = logpdf[i_sample] - log_diff = log_g[:, None] - log_g[None, :] - weighted_exp = w_s[:, None] * nx.exp(log_diff) - denom = nx.sum(weighted_exp, axis=0)[:, None] * nx.ones(plan.shape[1]) - p_mat = plan / denom + log_denom = logsumexp(log_g, scaling_factor=w_s) + p_mat = plan * nx.exp( + log_g.reshape((k_s, 1)) - log_denom + ) # shape (k_s, k_t): p_mat[i,j] = plan[i,j]*g_i(x)/D(x) p = p_mat.reshape(k_s * k_t) # stack line-by-line # sample between 0 and k_s * k_t - 1 From e8eb2e1aa7600a57c1a2b73cc60977e924eb0529 Mon Sep 17 00:00:00 2001 From: Jonathan Legrand Date: Mon, 21 Sep 2026 11:50:56 +0200 Subject: [PATCH 3/7] Handle the no scaling case --- ot/gmm.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/ot/gmm.py b/ot/gmm.py index 7c2c42e7c..65468ff0b 100644 --- a/ot/gmm.py +++ b/ot/gmm.py @@ -251,9 +251,28 @@ def gmm_ot_plan(m_s, m_t, C_s, C_t, w_s, w_t, log=False): def logsumexp(x, scaling_factor=None): """ - logsumexp trick such as in https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/ + Computes log(sum(scaling_factor * exp(x))) stably using the log-sum-exp trick + with optional per-element weight. + + Parameters + ---------- + x : array-like + Log-values to sum. + scaling_factor : array-like, optional + Weights for each term. If None, defaults to 1 + + Returns + ------- + float + log(sum(scaling_factor * exp(x))), computed stably. + + References + ---------- + Gundersen, G. (2020). The Log-Sum-Exp trick. Blog Post. Retrieved from https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/ """ nx = get_backend(x, scaling_factor) + if scaling_factor is None: + scaling_factor = 1 shift = nx.max(x) y = shift + nx.log(nx.sum(scaling_factor * nx.exp(x - shift))) return y From c7f2f668c0a18ed86f2c6568c57474fac0906a32 Mon Sep 17 00:00:00 2001 From: Jonathan Legrand Date: Mon, 21 Sep 2026 11:54:02 +0200 Subject: [PATCH 4/7] Remove comment --- test/test_gmm.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/test/test_gmm.py b/test/test_gmm.py index 5e5c928ca..b1ab3a127 100644 --- a/test/test_gmm.py +++ b/test/test_gmm.py @@ -162,12 +162,6 @@ def test_gmm_apply_map(): def test_gmm_apply_map_overflow(nx): - # In the original implementation - # log_diff computes log(g_i) - log(g_j) = log(g_i/g_j) - # so that g_k / sum_i(g_i) is easier to compute afterwards - # Problem is that if g_j(x) is small and g_i(x) is bigger, - # we can end up with high log ratios which overflows the - # subsequent exp calculation and throws a very unpleasant warning x_coord = 12 d = 1 k = 2 From 174d032a256f0007d8b90404daedf94a43bdd962 Mon Sep 17 00:00:00 2001 From: Jonathan Legrand Date: Mon, 21 Sep 2026 14:01:34 +0200 Subject: [PATCH 5/7] Require scaling factor in logsumexp --- ot/gmm.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/ot/gmm.py b/ot/gmm.py index 65468ff0b..d39206d3f 100644 --- a/ot/gmm.py +++ b/ot/gmm.py @@ -249,32 +249,33 @@ def gmm_ot_plan(m_s, m_t, C_s, C_t, w_s, w_t, log=False): return emd(w_s, w_t, D, log=log) -def logsumexp(x, scaling_factor=None): +def logsumexp(a, scaling_factor): """ Computes log(sum(scaling_factor * exp(x))) stably using the log-sum-exp trick - with optional per-element weight. + with per-element weight. The backend nx.logsumexp does not allow passing + scaling weights. Parameters ---------- - x : array-like + a : array-like Log-values to sum. - scaling_factor : array-like, optional - Weights for each term. If None, defaults to 1 + scaling_factor : array-like + Weights for each term, must be of the same shape as a. Returns ------- float - log(sum(scaling_factor * exp(x))), computed stably. + log(sum(scaling_factor * exp(a))), computed stably. References ---------- Gundersen, G. (2020). The Log-Sum-Exp trick. Blog Post. Retrieved from https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/ """ - nx = get_backend(x, scaling_factor) + nx = get_backend(a, scaling_factor) if scaling_factor is None: scaling_factor = 1 - shift = nx.max(x) - y = shift + nx.log(nx.sum(scaling_factor * nx.exp(x - shift))) + shift = nx.max(a) + y = shift + nx.log(nx.sum(scaling_factor * nx.exp(a - shift))) return y @@ -363,8 +364,8 @@ def gmm_ot_apply_map( # i and j, b[i, j] is the translation part rng = np.random.RandomState(seed) - A = nx.zeros((k_s, k_t, d, d)) - b = nx.zeros((k_s, k_t, d)) + A = nx.zeros((k_s, k_t, d, d), type_as=C_s) + b = nx.zeros((k_s, k_t, d), type_as=A) # only need to compute for non-zero plan entries for i, j in zip(*nx.where(plan > 0)): From 99c75562454535afa54cecd426f139c189b8fd89 Mon Sep 17 00:00:00 2001 From: Jonathan Legrand Date: Mon, 21 Sep 2026 14:02:07 +0200 Subject: [PATCH 6/7] Test logsumexp for torch and jax backends --- test/test_gmm.py | 34 ++++++++++++++++++++-------------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/test/test_gmm.py b/test/test_gmm.py index b1ab3a127..bd9324343 100644 --- a/test/test_gmm.py +++ b/test/test_gmm.py @@ -19,6 +19,7 @@ gmm_ot_apply_map, gmm_ot_plan_density, gmm_barycenter_fixed_point, + logsumexp, ) try: @@ -161,26 +162,22 @@ def test_gmm_apply_map(): gmm_ot_apply_map(x, m_s, m_t, C_s, C_t, w_s, w_t, plan=plan) +@pytest.skip_backend("tf") # skips because of array assignment +@pytest.skip_backend("jax") def test_gmm_apply_map_overflow(nx): - x_coord = 12 + x_coord = 12.0 d = 1 k = 2 - x = np.array([x_coord]).reshape((-1, d)) + x = nx.from_numpy(np.array([x_coord]).reshape((-1, d))) - m_s = np.array([x_coord, 0]).reshape((k, d)) - m_t = m_s.copy() + m_s = nx.from_numpy(np.array([x_coord, 0.0], dtype=np.float64).reshape((k, d))) + m_t = nx.from_numpy(np.array([x_coord, 0.0], dtype=np.float64).reshape((k, d))) - C_s = np.ones(m_s.shape).reshape((k, d, d)) / 10 - C_t = C_s.copy() + C_s = nx.from_numpy(np.ones((k, d, d)) / 10.0) + C_t = nx.from_numpy(np.ones((k, d, d)) / 10.0) - w_s = np.array([0.5, 0.5]) - w_t = w_s.copy() - m_s = nx.from_numpy(m_s) - m_t = nx.from_numpy(m_t) - C_s = nx.from_numpy(C_s) - C_t = nx.from_numpy(C_t) - w_s = nx.from_numpy(w_s) - w_t = nx.from_numpy(w_t) + w_s = nx.from_numpy(np.array([0.5, 0.5])) + w_t = nx.from_numpy(np.array([0.5, 0.5])) with warnings.catch_warnings(): warnings.simplefilter(action="error") @@ -197,6 +194,15 @@ def test_gmm_apply_map_overflow(nx): ) +def test_logsumexp_overflow_safety(nx): + """Test that large exponents don't overflow.""" + x = nx.from_numpy(np.array([0.0, 710.0])) + result = logsumexp(x, scaling_factor=nx.from_numpy(np.array([1, 1]))) + # Should be equal to exp(log(0)) + exp(log(710)) = 710 + assert nx.isfinite(result) + assert nx.allclose(result, nx.from_numpy(np.array([710.0]))) + + @pytest.mark.skipif(not torch, reason="No torch available") def test_gradient_gmm_ot_loss_pytorch(): m_s, m_t, C_s, C_t, w_s, w_t = get_gmms() From a63c23643ed70eec18fa464f13c27258c6b55cbc Mon Sep 17 00:00:00 2001 From: Jonathan Legrand Date: Thu, 24 Sep 2026 15:18:55 +0200 Subject: [PATCH 7/7] Compute T_mean with logsumexp --- ot/gmm.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/ot/gmm.py b/ot/gmm.py index d39206d3f..778f83b2a 100644 --- a/ot/gmm.py +++ b/ot/gmm.py @@ -249,7 +249,7 @@ def gmm_ot_plan(m_s, m_t, C_s, C_t, w_s, w_t, log=False): return emd(w_s, w_t, D, log=log) -def logsumexp(a, scaling_factor): +def logsumexp(a, scaling_factor, axis=None): """ Computes log(sum(scaling_factor * exp(x))) stably using the log-sum-exp trick with per-element weight. The backend nx.logsumexp does not allow passing @@ -275,7 +275,7 @@ def logsumexp(a, scaling_factor): if scaling_factor is None: scaling_factor = 1 shift = nx.max(a) - y = shift + nx.log(nx.sum(scaling_factor * nx.exp(a - shift))) + y = shift + nx.log(nx.sum(scaling_factor * nx.exp(a - shift), axis=axis)) return y @@ -352,10 +352,14 @@ def gmm_ot_apply_map( # gaussian mapping between components i and j applied to x T_ij_x = x @ A + b - z = w_s[:, None, None] * nx.exp(logpdf - logpdf[i][None, :, :]) - denom = nx.sum(z, axis=0) - out = out + plan[i, j] * T_ij_x / denom + log_g_i_x = logpdf[i] + # Could be optimized, that's not too smart to compute denom here at each iteration + denom = logsumexp( + logpdf.squeeze(), scaling_factor=w_s.reshape((-2, 1)), axis=0 + ) + p_ij_x = plan[i, j] * nx.exp(log_g_i_x - denom.reshape((-2, 1))) + out = out + p_ij_x * T_ij_x return out