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
• Update to ott-jax 0.5.0 #758

[MUCDK]

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

[selmanozleyen]

Hi @MUCDK, except the notebooks I think it is ready to merge. But here are noteworthy side-effects that was caused by the ottjax update to 0.5.0

• inner_iterations parameter removed since it was fixed to 1 in new version of ottjax
• we no longer take str for initializers since ottjax dropped the support for this. We take initializer objects. I modified the tests accordingly but notebooks also need to comply to this

[MUCDK]

Can we keep strings for initializers , or do Literal[...], init_lib.SinkhornInitializer and port the code to src?

[MUCDK]

This looks great, thanks @selmanozleyen !

Wdyt about allowing strings for initialisers additionally? Might be too hard for people not familiar with these things otherwise.

[selmanozleyen]

This looks great, thanks @selmanozleyen !

Wdyt about allowing strings for initialisers additionally? Might be too hard for people not familiar with these things otherwise.

The issue there is differentiating the arguments required when creating the initializer objects vs the arguments required when calling initializers. For example you don't know if rank belongs to RandomInitializer or solve(rank=rank) in initializer_kwargs unless we have a hard coded list. We might not be able to use inspect module because some classes might be using kwargs in their constructors (because of this we won't know what the intializer takes). I will look if we can use inspect module but I think providing a class for creating initializers from strings is a better approach to clarify which argument goes where.

[MUCDK]

This looks great, thanks @selmanozleyen !
Wdyt about allowing strings for initialisers additionally? Might be too hard for people not familiar with these things otherwise.

The issue there is differentiating the arguments required when creating the initializer objects vs the arguments required when calling initializers. For example you don't know if rank belongs to RandomInitializer or solve(rank=rank) in initializer_kwargs unless we have a hard coded list. We might not be able to use inspect module because some classes might be using kwargs in their constructors (because of this we won't know what the intializer takes). I will look if we can use inspect module but I think providing a class for creating initializers from strings is a better approach to clarify which argument goes where.

I see. Yes, then maybe let's create a class , but then call it from inside the solve method, because moscot users are not expected to know OOP. We are really supposed to do it in a simple manner.

[selmanozleyen]

hi @MUCDK , I made the changes. If the user wants to create more advanced initializers (ie giving something to the initializer constructor as kwargs) they can do by creating it themselves. I think its reasonable to assume that someone who'd want to do that they can look ott initializers up.

[MUCDK]

That looks great, thanks so much @selmanozleyen !

[MUCDK]

Please merge, and then do a new release v0.4.0 of moscot!