Skip to content

Batch OT losses (Sinkhorn + Gromov) #755

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

Merged
rflamary merged 37 commits into from
Sep 16, 2025
Merged

Batch OT losses (Sinkhorn + Gromov) #755

Changes from 1 commit
Commits
Show all changes
37 commits
Select commit Hold shift + click to select a range
2bf1a0f
linear ot implemented
Jul 23, 2025
b24c275
improve stopping criterion and assymetric case
Jul 24, 2025
1267322
Add recompute_const and simplify the pipeline for the symmetric = False
Aug 22, 2025
acb6233
add tests
Aug 23, 2025
8645bb8
update the examples and rename to follow the "ot.solve" naming conven…
Aug 25, 2025
ba09929
update realeases.md
Aug 25, 2025
a3344cb
idem
Aug 25, 2025
f5dfe15
Merge branch 'master' into ot-batch
rflamary Aug 25, 2025
022e398
move set_grad_enabled to backend
Aug 25, 2025
aee7eb4
set_grad_enabled for quadratric solver
Aug 25, 2025
fcff735
update doc
Aug 25, 2025
9ccb357
remove useless importation in doc
Aug 25, 2025
e04d4e0
Update references
Aug 25, 2025
68117dd
Merge branch 'master' into ot-batch
cedricvincentcuaz Aug 27, 2025
ed8c842
Merge branch 'master' into ot-batch
rflamary Sep 4, 2025
fe37b7e
Merge branch 'master' into ot-batch
rflamary Sep 8, 2025
ae478fe
update example
Sep 10, 2025
a8e2319
Remove classes in quadratic, move examples to backend, add potentials…
Sep 10, 2025
395e4cc
updat tests
Sep 10, 2025
0d6fbbc
Massive improvement of the documentation for ot.batch
Sep 10, 2025
da7639e
cover (almost) all ot.batch with tests
Sep 11, 2025
62a6f9d
bug in the tests
Sep 11, 2025
17fa680
update docstring
Sep 11, 2025
6f6109d
highlight that ot.batch is solving the entropic version
Sep 11, 2025
5524efb
removing yet another error in the docstring
Sep 11, 2025
bb37cf2
Add missing parameter recompute_const
Sep 11, 2025
f6ed5f6
Remove png, add all backends and gradient mode to tests
Sep 15, 2025
cde102b
add the missing pytest
Sep 15, 2025
f06b99f
change .sum() into nx.sum
Sep 15, 2025
dd6eda3
add missing backend
Sep 15, 2025
3fa7c1b
yet another missing nx
Sep 15, 2025
a77d034
remove useless squeeze and add test for non-log bregman
Sep 15, 2025
1528f95
remove last_step from quadratic tests
Sep 15, 2025
69f59a1
add missing tests and improve documentation
Sep 16, 2025
27cea5e
proper unsqueeze test
Sep 16, 2025
3264b5c
add unsqueeze to tensorflow
Sep 16, 2025
e70ecd1
solve double backprop issue in test_gradients_torch
Sep 16, 2025
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
Prev Previous commit
Next Next commit
set_grad_enabled for quadratric solver
  • Loading branch information
PaulKrzakala committed Aug 25, 2025
commit aee7eb459cc4da56f604b9aec0b4e41411118733
4 changes: 2 additions & 2 deletions ot/batch/_quadratic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
from ..utils import OTResult
from ot.backend import get_backend
from ot.batch._linear import loss_linear_batch
from ot.batch._utils import grad_enabled, bmv, bop, bregman_log_batch
from ot.batch._utils import bmv, bop, bregman_log_batch
from abc import ABC, abstractmethod


Expand Down Expand Up @@ -447,7 +447,7 @@ def solve_gromov_batch(

grad_inner = "autodiff" if grad == "autodiff" else "detach"

with grad_enabled(nx, grad == "autodiff"):
with nx.set_grad_enabled(grad == "autodiff"):
T = T_init
log_T = nx.log(T + 1e-15) # Avoid log(0)
for iter in range(max_iter):
Expand Down
Loading