diff --git a/RELEASES.md b/RELEASES.md index 063701229..2e272e184 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -15,6 +15,7 @@ #### Closed issues +- Fix `ot.bregman.screenkhorn` failing with a shape error when `ns_budget` and `nt_budget` are the full source and target sizes and these differ (PR #871) - 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) - 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) diff --git a/ot/bregman/_screenkhorn.py b/ot/bregman/_screenkhorn.py index ea00a03cc..b562c0188 100644 --- a/ot/bregman/_screenkhorn.py +++ b/ot/bregman/_screenkhorn.py @@ -182,8 +182,8 @@ def projection(u, epsilon): K_IJc = [] K_IcJ = [] - vec_eps_IJc = nx.zeros((nt,), type_as=M) - vec_eps_IcJ = nx.zeros((ns,), type_as=M) + vec_eps_IJc = nx.zeros((ns,), type_as=M) + vec_eps_IcJ = nx.zeros((nt,), type_as=M) else: # sum of rows and columns of K diff --git a/test/test_bregman.py b/test/test_bregman.py index 17b400306..ddfefad3b 100644 --- a/test/test_bregman.py +++ b/test/test_bregman.py @@ -1434,6 +1434,20 @@ def test_screenkhorn(nx): np.testing.assert_allclose(G_sink.sum(1), G_screen.sum(1), atol=1e-02) + +def test_screenkhorn_full_budget_non_square(): + # with the full budget no point is screened out, which used to fail with + # a shape error whenever the source and target sizes differ + rng = np.random.RandomState(0) + ns, nt = 20, 12 + a = ot.unif(ns) + b = ot.unif(nt) + M = ot.dist(rng.randn(ns, 2), rng.randn(nt, 2)) + + G_screen = ot.bregman.screenkhorn(a, b, M, 1.0, ns_budget=ns, nt_budget=nt) + np.testing.assert_allclose(G_screen, ot.sinkhorn(a, b, M, 1.0), atol=1e-4) + + def test_convolutional_barycenter_non_square(nx): # test for image with height not equal width A = np.ones((2, 2, 3)) / (2 * 3)