Bugs in the DualPotentials class

[lucaeyring]

I believe there are two bugs in the DualPotentials class:

1. in the distance function tgt and src are swapped. g should be applied to the src, not the tgt.
2. It is not possible to move a DualPotentials object to a different device because its inputs f and g are Callables which is not a valid Jax type.

To Reproduce
Steps to reproduce the behavior of 2:

import jax
from ott.problems.linear import potentials

f = lambda x: x
g = lambda x: x
test_potential = potentials.DualPotentials(f, g, corr=True)
jax.device_put(test_potential, "cpu")

Expected behavior
One should be able to move a DualPotentials object to a different device.

A possible fix for 2 would be to instead use the TrainStates as input.

[michalk8]

1. Do you mean when corr=False or corr=True or in both cases? Will take a look again at the Makkuva et al. paper.
2. Thanks, nice catch! Will be fixed in Fix jax.device_put for potentials #183

[michalk8]

You're right, seems that it's swapped in the paper (also not consistent with NeuralDual's way of computing the W2 distance. Will be addressed in a different PR than #183

[marcocuturi]

@michalk8 , shouldnt we fix this?

[bamos]

@michalk8 , shouldnt we fix this?

@marcocuturi I have a fix swapping the order of the neural dual potentials to be consistent with Makkuva et al. in a larger PR I'll send soon, but I've still been having some issues stabilizing it and setting up a version with a minimal amount of changes to start a larger discussion. My updates unfortunately ended up being more difficult to clean up and debug than I expected, but I'll send an initial version in a few days

[marcocuturi]

@michalk8 , no rush at all, let's chat live about this then!!

[michalk8]

closed via #219