Batch OT losses (Sinkhorn + Gromov)#755
Merged
rflamary merged 37 commits intoPythonOT:masterfrom Sep 16, 2025
Merged
Conversation
added 7 commits
July 23, 2025 18:03
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #755 +/- ##
========================================
Coverage 97.15% 97.16%
========================================
Files 101 107 +6
Lines 21386 21879 +493
========================================
+ Hits 20778 21259 +481
- Misses 608 620 +12 🚀 New features to boost your workflow:
|
rflamary
changed the title
rflamary
reviewed
Sep 4, 2025
Collaborator
rflamary
left a comment
rflamary
left a comment
There was a problem hiding this comment.
Thanks @KrzakalaPaul here are a few comments to take into account please.
ot/batch/_linear.py
Outdated
| log_dual=True, | ||
| grad="detach", | ||
| ): | ||
| r"""Solves the linear optimal transport problem using Bregman projections. |
Collaborator
There was a problem hiding this comment.
Add entropic here this is not exact OT
ot/batch/_quadratic.py
Outdated
| symmetric=None, | ||
| M=None, | ||
| alpha=None, | ||
| epsilon=1e-2, |
…, remove context managers for grads. To do: improve doc and tests
rflamary
reviewed
Sep 12, 2025
ot/batch/_quadratic.py
Outdated
| tol=1e-5, | ||
| max_iter_inner=50, | ||
| tol_inner=1e-5, | ||
| grad="detach", |
Collaborator
There was a problem hiding this comment.
we need a gradient by default
added 11 commits
September 15, 2025 15:10
MatDag
pushed a commit
to MatDag/POT
that referenced
this pull request
Sep 25, 2025
* linear ot implemented * improve stopping criterion and assymetric case * Add recompute_const and simplify the pipeline for the symmetric = False * add tests * update the examples and rename to follow the "ot.solve" naming conventions * update realeases.md * idem * move set_grad_enabled to backend * set_grad_enabled for quadratric solver * update doc * remove useless importation in doc * Update references * update example * Remove classes in quadratic, move examples to backend, add potentials, remove context managers for grads. To do: improve doc and tests * updat tests * Massive improvement of the documentation for ot.batch * cover (almost) all ot.batch with tests * bug in the tests * update docstring * highlight that ot.batch is solving the entropic version * removing yet another error in the docstring * Add missing parameter recompute_const * Remove png, add all backends and gradient mode to tests * add the missing pytest * change .sum() into nx.sum * add missing backend * yet another missing nx * remove useless squeeze and add test for non-log bregman * remove last_step from quadratic tests * add missing tests and improve documentation * proper unsqueeze test * add unsqueeze to tensorflow * solve double backprop issue in test_gradients_torch --------- Co-authored-by: PaulKrzakala <paul.krzakala@gmail.com> Co-authored-by: Rémi Flamary <remi.flamary@gmail.com> Co-authored-by: Cédric Vincent-Cuaz <cedvincentcuaz@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Types of changes
Add the ot.batch module for solving N optimal transport problems at the same time in parallel. The two main features are:
Some examples have been added to examples/batch
Motivation and context / Related issue
It is often the case that one needs to solve N optimal transport problem at the same time. At the moment, the only way to do this in POT was using a for loop. This is very inefficient for data stored on a GPU for instance.
Instead, it is now possible to use ot.batch for solving N problems at the same time with fully batch parallel operations. As demonstrated in examples/batch/demo_efficiency.py this can lead to a speed of 1 or 2 orders of magnitude.
How has this been tested (if it applies)
New tests have been added to test/batch/ to ensure that using ot.solve_batch yields the same results that using a for loops of ot.solve. Idem for ot.solve_batch_gromov.
PR checklist