diff --git a/RELEASES.md b/RELEASES.md index a83c8ba2d..dc0c9478e 100644 --- a/RELEASES.md +++ b/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 `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) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index ad27c6e91..90482afce 100644 --- a/ot/solvers/_linear.py +++ b/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( diff --git a/test/test_solvers.py b/test/test_solvers.py index 6cc8137cc..1d7e59403 100644 --- a/test/test_solvers.py +++ b/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):