Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor/handle solve args #748

Open
wants to merge 26 commits into
base: main
Choose a base branch
from
Open

Conversation

selmanozleyen
Copy link
Collaborator

@selmanozleyen selmanozleyen commented Sep 22, 2024

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:

@selmanozleyen selmanozleyen marked this pull request as draft September 22, 2024 21:35
@@ -435,7 +435,7 @@ def _prepare(
cost_matrix_rank: Optional[int] = None,
time_scales_heat_kernel: Optional[TimeScalesHeatKernel] = None,
# problem
alpha: float = 0.5,
alpha: Optional[float] = None,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that makes sense

Copy link
Collaborator

@MUCDK MUCDK left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
2 participants