[MRG] Compute value in ot.solve_sample with reg and lazy=True

RELEASES.md

@@ -12,6 +12,7 @@
 
 #### Closed issues
 
+- Fix `ot.solve_sample` with `reg` and `lazy=True` returning `value=None`, which also made the debiased lazy Sinkhorn divergence (`debias=True` or `"split"`) fail with a `TypeError` (PR #861)
 - 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)
 - 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)

ot/solvers/_linear.py

@@ -1569,6 +1569,17 @@ def solve_sample(
 
         else:
             if unbalanced is None:
+                if not isinstance(reg_type, str) or reg_type.lower() not in [
+                    "kl",
+                    "entropy",
+                ]:
+                    raise (
+                        NotImplementedError(
+                            'Not implemented reg_type="{}" with lazy=True'.format(
+                                reg_type
+                            )
+                        )
+                    )
                 if max_iter is None:
                     max_iter = 1000
                 if tol is None:
@@ -1594,6 +1605,23 @@ def solve_sample(
                 potentials = (log["u"], log["v"])
                 lazy_plan = log["lazy_plan"]
 
+                # regularized value, accumulated over batches of rows of the
+                # plan to keep the memory cost of the lazy solver
+                if a is None:
+                    a = nx.ones(X_a.shape[0], type_as=X_a) / X_a.shape[0]
+                if b is None:
+                    b = nx.ones(X_b.shape[0], type_as=X_b) / X_b.shape[0]
+                reg_value = 0
+                for i in range(0, X_a.shape[0], batch_size):
+                    plan_i = lazy_plan[i : i + batch_size]
+                    if reg_type.lower() == "entropy":
+                        reg_value += nx.sum(plan_i * nx.log(plan_i + 1e-16))
+                    else:
+                        reg_value += nx.kl_div(
+                            plan_i, a[i : i + batch_size, None] * b[None, :] + 1e-15
+                        )
+                value = value_linear + reg * reg_value
+
             else:
                 raise (
                     NotImplementedError(

test/test_solvers.py

@@ -631,6 +631,41 @@ def test_solve_sample_lazy(nx):
 
     np.testing.assert_allclose(sol0.plan, sol.lazy_plan[:], rtol=1e-5, atol=1e-5)
 
+@pytest.mark.parametrize("reg_type", ["KL", "entropy"])
+def test_solve_sample_lazy_value(nx, reg_type):
+    rng = np.random.RandomState(0)
+    X_s = rng.randn(25, 2)
+    X_t = rng.randn(15, 2) + 1
+    a = rng.rand(25)
+    a /= a.sum()
+    b = ot.utils.unif(15)
+
+    X_s, X_t, a, b = nx.from_numpy(X_s, X_t, a, b)
+
+    sol0 = ot.solve_sample(X_s, X_t, a, b, reg=1, reg_type=reg_type)
+    # small batches so that the value is accumulated over several of them
+    sol = ot.solve_sample(
+        X_s, X_t, a, b, reg=1, reg_type=reg_type, lazy=True, batch_size=4
+    )
+    np.testing.assert_allclose(
+        nx.to_numpy(sol0.value), nx.to_numpy(sol.value), rtol=1e-5
+    )
+
+    # debiased values combine the values of several lazy problems
+    for debias in [True, "split"]:
+        sol0 = ot.solve_sample(
+            X_s, X_t, a, b, reg=1, reg_type=reg_type, debias=debias
+        )
+        sol = ot.solve_sample(
+            X_s, X_t, a, b, reg=1, reg_type=reg_type, debias=debias, lazy=True
+        )
+        np.testing.assert_allclose(
+            nx.to_numpy(sol0.value), nx.to_numpy(sol.value), rtol=1e-5
+        )
+
+    with pytest.raises(NotImplementedError):
+        ot.solve_sample(X_s, X_t, a, b, reg=1, reg_type="L2", lazy=True)
+
 
 @pytest.mark.parametrize("metric", ["sqeuclidean", "euclidean", "cityblock"])
 def test_solve_sample_lazy_emd(nx, metric):