Refactor/handle solve args

[selmanozleyen]

hi @MUCDK ,

So good news is we currently do a good job on partitioning the kwargs for solve. In solve we give any kwarg we don't know to either SinkhornSolver or GWSolver constructors. SinkhornSolver uses Sinkhorn or LRSinkhorn from ottjax, these classes don't have kwargs in their constructors so when using SinkhornSolver as a backend we are good. GWSolver uses GromovWasserstein or LRGromovWasserstein from ottjax. The parent class of these class WassersteinSolver don't throw an error on unrecognized args. The tests will pass after the ottjax PR merges.

Here is the PR in ott-jax: ott-jax/ott#579

Other things done:

• Added tests if we throw appropriate errors on completely unrecognized arguments
• I refactored where we handle checks for alpha, now it's completely independent from CompoundProblem or any other more abstract class. It's handled in GWSolver as it should.
• I added extra tests on the errors we raise for alpha or the data given.

Additionally closes:

• turns out we can remove the skip for these test so I will close No convergence with new commit of ottjax in LRGW solvers unless inner_iterations=10 is set #678
• check the convergence flag #703
• check for unrecognized in solve methods #720

[MUCDK]

can we just make a comment behind this that default is 0.5

[MUCDK]

I'm actually wondering whether we should have a default alpha whenever we don't want it to be 1.0

I.e. always set it explicitly in the classes which use GW, wdyt?

[selmanozleyen]

I set it like this so that it doesn't change the current behaviour.

I.e. always set it explicitly in the classes which use GW, wdyt?

You mean to make non-optional? I personally prefer non-optional parameters, especially if the class is an internal solver. I also think we should rename GWSolver to FGWSolver (because it technically can solve fgw and gw) and just set alpha=1 when in a GWProblem.

[MUCDK]

I think that makes sense

[MUCDK]

Thanks, overall really nice. Let's see what Mike and Marco say on the ott-jax side!