Skip to content

[MRG] Fix screenkhorn with the full budget on non-square problems #871

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
itzzdev09 wants to merge 2 commits into
base: master
Choose a base branch
from
+17 −2
Open

[MRG] Fix screenkhorn with the full budget on non-square problems #871

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6cb7106
Fix screenkhorn with the full budget on non-square problems
itzzdev09 Sep 21, 2026
d8c942c
Loosen the screenkhorn tolerance and add the PR number to the release…
itzzdev09 Sep 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions RELEASES.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions ot/bregman/_screenkhorn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 14 additions & 0 deletions test/test_bregman.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading