[MRG] Sinkhorn gradient last step by smazelet · Pull Request #693 · PythonOT/POT
Types of changes
I added a grad=last_step feature to the ot.solvers.solve function to limit the gradient computation to the last iteration of Sinkhorn instead of computing the gradient for every iteration.
Motivation and context / Related issue
This change is required in the case where computing the gradient for all the Sinkhorn iterations is too heavy.
How has this been tested (if it applies)
I created a test_solve_last_step test to check that the gradient is different when using grad=last_step and grad=autodiff.
PR checklist
- I have read the CONTRIBUTING document.
- The documentation is up-to-date with the changes I made (check build artifacts).
- All tests passed, and additional code has been covered with new tests.
- I have added the PR and Issue fix to the RELEASES.md file.