[MRG] Fix GMM rand map overflow

ot/gmm.py

@@ -249,6 +249,36 @@ 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, 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
+    scaling weights.
+
+    Parameters
+    ----------
+    a : array-like
+        Log-values to sum.
+    scaling_factor : array-like
+        Weights for each term, must be of the same shape as a.
+
+    Returns
+    -------
+    float
+        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(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), axis=axis))
+    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
 ):
@@ -322,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
 
@@ -334,8 +368,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)):
@@ -354,10 +388,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

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
@@ -18,6 +19,7 @@
     gmm_ot_apply_map,
     gmm_ot_plan_density,
     gmm_barycenter_fixed_point,
+    logsumexp,
 )
 
 try:
@@ -160,6 +162,47 @@ 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.0
+    d = 1
+    k = 2
+    x = nx.from_numpy(np.array([x_coord]).reshape((-1, d)))
+
+    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 = nx.from_numpy(np.ones((k, d, d)) / 10.0)
+    C_t = nx.from_numpy(np.ones((k, d, d)) / 10.0)
+
+    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")
+        gmm_ot_apply_map(
+            x,
+            m_s,
+            m_t,
+            C_s,
+            C_t,
+            w_s,
+            w_t,
+            method="rand",
+            seed=0,
+        )
+
+
+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()