Fix jax.device_put for potentials

[michalk8]

Fixes jax.device_put on DualPotentials.

partially addresses #182

[codecov-commenter]

Codecov Report

Merging #183 (9b28c55) into main (156ef2b) will decrease coverage by 0.04%.
The diff coverage is 64.70%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #183      +/-   ##
==========================================
- Coverage   89.28%   89.23%   -0.05%     
==========================================
  Files          51       51              
  Lines        5187     5192       +5     
  Branches      527      529       +2     
==========================================
+ Hits         4631     4633       +2     
- Misses        430      431       +1     
- Partials      126      128       +2     

Impacted Files	Coverage Δ
ott/problems/linear/potentials.py	84.76% <64.70%> (-2.24%)	⬇️

[marcocuturi]

That's great Michal!

There was a change in the last PR that I overlooked on passing a and b by default. Because f and g here correspond exactly to an existing a and b (even if that pair was defined by default initially, it somewhat becomes unique once it is set), it would be a mistake IMHO to allow the user to pass on an f and a g without requiring them to pass the corresponding a and b. This could easily lead to a silent error where the f,g were computed for a given pair of weights but the dual potentials evaluated with the default uniform measures. Therefore I would prefer sticking to the way this was handled before (see also this comment: #167 (comment))

[michalk8]

That's great Michal!

There was a change in the last PR that I overlooked on passing a and b by default. Because f and g here correspond exactly to an existing a and b (even if that pair was defined by default initially, it somewhat becomes unique once it is set), it would be a mistake IMHO to allow the user to pass on an f and a g without requiring them to pass the corresponding a and b. This could easily lead to a silent error where the f,g were computed for a given pair of weights but the dual potentials evaluated with the default uniform measures. Therefore I would prefer sticking to the way this was handled before (see also this comment: #167 (comment))

I'd then consider passing LinearProblem instead of geom/a/b to further avoid any user-level mistakes.