[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.