[MRG] Add set_gradients method for JAX backend.

[AdrienCorenflos]

Types of changes

• Docs change / refactoring / dependency upgrade
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

Motivation and context / Related issue

The set_gradient is possible in JAX.

#277

How has this been tested (if it applies)

Added a modified unittest for JAX.

Checklist

• [ X ] The documentation is up-to-date with the changes I made.
• [ X ] I have read the CONTRIBUTING document.
• [ X ] All tests passed, and additional code has been covered with new tests.

[codecov]

Codecov Report

Merging #278 (eb8d0cf) into master (14c30d4) will increase coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #278      +/-   ##
==========================================
+ Coverage   92.64%   92.66%   +0.01%     
==========================================
  Files          19       19              
  Lines        3754     3761       +7     
==========================================
+ Hits         3478     3485       +7     
  Misses        276      276              

[rflamary]

This is awesome! that's why you need people who know the framework for multiple backend.

Could you please also check the you can compute gradients for emd2? by updating this test?
https://github.com/PythonOT/POT/blob/master/test/test_ot.py#L84

[AdrienCorenflos]

Well, no you won't be able to. You are casting your tensors to numpy, and JAX tracing is not compatible with this... In order to use JAX you would need to dispatch the call to the host by using a callback instead (see https://jax.readthedocs.io/en/latest/jax.experimental.host_callback.html), and only after this set the gradients. In practice this means that the backends should essentially have a kind of pycall method.

[AdrienCorenflos]

practice this means that the backends should essentially have a kind of pycall method.

Note that this is feasible, but it's a bit more work than simply setting gradients, and requires some preliminary design thinking.

[rflamary]

OK i'm sorry i vnever understoiod the subtleties of jax. But if what you implemented did not define the gradient for a given variable, what is its use?

[AdrienCorenflos]

It does replace the gradient within JAX code, but it does not allow to bypass the tracing. The problem of emd2 is not with the gradient, it's with the fact that you have an operation that JAX can't trace in the middle.

[rflamary]

Note that We use the set_gradient in emd2 exactly to bypass the need for numpy arrays and it works for torch. This is why i tried to define a new function with @custom_jvp (and call it) in set gradients but i wasn't a good strategy.