Enable to set parameters for inner Sinkhorn loop in Gromov-Wasserstein

[ott-jax]

Enable to set parameters for inner Sinkhorn loop in Gromov-Wasserstein

[marcocuturi]

Thanks Dominik for the great comment.

I think this is possible using the most recent API, as illustrated in this colab:

https://colab.research.google.com/drive/1YWKKjSWnpbhpeMae2tlkHKLLZXP97Cij?usp=sharing

You need to instantiate manually a QuadraticProblem, use a GromovWasserstein solver on it, and specify in GromovWasserstein the linear_ot_solver (this is common for all WassersteinSolver.
Namely, the linear solver inside the GW solver can be provided explicitly, with all parameters you might want, which is what happens in that call in the colab

GWsolver= gw.GromovWasserstein(
    linear_ot_solver=sinkhorn.Sinkhorn(
      min_iterations=1, 
      max_iterations=30,
      threshold=1e-1),
    min_iterations=1,
    max_iterations=7,
    epsilon=100.0,
    )

and then simply make the call

quad_p = qp.QuadraticProblem(geom_xx, geom_yy, scale_cost=True)
out = GWsolver(quad_p)

In the colab you can see how the .errors field of the output of GW shows the right size, showing how max_iterations are properly handled.

Using such classes, it is easy to have access to detailed granularity on the components that make up an OT problem. I think it's also a lot more viable than overloading the signature of gromov_wasserstein with sinkhorn_kwargs or sinkhorn_lr__kwargs etc... This is also why @michalk8 wants to get rid of make functions #131

[MUCDK]

Thanks Marco!
I like the idea with sinkhorn_kwargs but the current approach is also nice!
Just want to make sure you are not about to refactor this?

[marcocuturi]

yes, the example I provided is clearly te way Michal and I want to go. We'll still maintain the sinkhorn and gromov_wasserstein wrappers for a little while, but I think we'll deprecate them at some point!

[MUCDK]

Amazing, thank you!