Skip to content
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.

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

add MMSinkhorn solver #549

Merged
merged 26 commits into from
Jun 21, 2024
Merged

add MMSinkhorn solver #549

merged 26 commits into from
Jun 21, 2024

Conversation

marcocuturi
Copy link
Contributor

@marcocuturi marcocuturi commented Jun 18, 2024

experimental code to compute the multimarginal Sinkhorn algorithm, as described in Algorithm 1 of https://arxiv.org/abs/2405.19532v1

Largely based on the Sinkhorn template, but simpler (no definition of MM problems, not unbalanced).

@marcocuturi marcocuturi added the enhancement New feature or request label Jun 19, 2024
@marcocuturi marcocuturi requested a review from michalk8 June 19, 2024 08:23
Copy link

codecov bot commented Jun 19, 2024

Codecov Report

Attention: Patch coverage is 93.08176% with 11 lines in your changes missing coverage. Please review.

Project coverage is 91.38%. Comparing base (84a1f1b) to head (79be6cd).
Report is 40 commits behind head on main.

Files with missing lines Patch % Lines
src/ott/experimental/mmsinkhorn.py 93.08% 9 Missing and 2 partials ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main     #549      +/-   ##
==========================================
+ Coverage   91.34%   91.38%   +0.03%     
==========================================
  Files          68       69       +1     
  Lines        7083     7242     +159     
  Branches      991     1019      +28     
==========================================
+ Hits         6470     6618     +148     
- Misses        463      472       +9     
- Partials      150      152       +2     
Files with missing lines Coverage Δ
src/ott/math/utils.py 91.11% <ø> (ø)
src/ott/solvers/linear/sinkhorn.py 99.38% <ø> (ø)
src/ott/experimental/mmsinkhorn.py 93.08% <93.08%> (ø)

Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

docs/experimental/mmsinkhorn.rst Outdated Show resolved Hide resolved
docs/experimental/mmsinkhorn.rst Show resolved Hide resolved
docs/references.bib Show resolved Hide resolved
src/ott/experimental/mmsinkhorn.py Show resolved Hide resolved
@michalk8 michalk8 self-requested a review June 20, 2024 12:31
docs/experimental/mmsinkhorn.rst Outdated Show resolved Hide resolved
docs/experimental/mmsinkhorn.rst Show resolved Hide resolved
src/ott/experimental/mmsinkhorn.py Outdated Show resolved Hide resolved
src/ott/experimental/mmsinkhorn.py Show resolved Hide resolved
src/ott/experimental/mmsinkhorn.py Outdated Show resolved Hide resolved
src/ott/experimental/mmsinkhorn.py Outdated Show resolved Hide resolved
src/ott/experimental/mmsinkhorn.py Outdated Show resolved Hide resolved
src/ott/experimental/mmsinkhorn.py Outdated Show resolved Hide resolved
src/ott/experimental/mmsinkhorn.py Outdated Show resolved Hide resolved
docs/experimental/mmsinkhorn.rst Show resolved Hide resolved
@marcocuturi marcocuturi requested a review from michalk8 June 20, 2024 21:34
@michalk8
Copy link
Collaborator

Looks great @marcocuturi , merging now!

@michalk8 michalk8 merged commit 787d4a9 into main Jun 21, 2024
10 of 12 checks passed
@michalk8 michalk8 deleted the MMS branch June 21, 2024 12:06
michalk8 added a commit that referenced this pull request Jun 27, 2024
* add MMSinkhorn solver

* more tests and properties

* docs

* correct docs bug

* add docstrings + test multiple costs + handle danskin properly

* spelling

* switch back to point cloud

* Fix jax>=0.4.30 CUDA installation on the CI

* Increase mem for `test_sinkhorn_online_memory_jit`

* spell

* fix

* fix error tolerances

* add tutorial

* typo

* typos

* typo

* polymatching

* pydocs

* spin off MMSinkhorn from SInkhorn, add back `norm_error`

* typos

* typo

* typo

* incorporate feedback + differentiability

* typo

* Fix minor pet-peeves

* Link to classes in the docs

---------

Co-authored-by: Michal Klein <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants