Optionally return transport plans for the Sinkhorn loss

[arnaudvl]

First of all, thanks for the great package (incl. KeOps). I believe it could be very helpful for the user to be able to inspect the optimal transport plans found by the Sinkhorn iterations. This could e.g. be done by optionally computing them after the sinkhorn_loop routine following Eq.9 of Interpolating between Optimal Transport and MMD using Sinkhorn Divergences (simplified for batch size of 1):

p_xy = (((1 / ε) * (f_ba.unsqueeze(-1) + g_ab.unsqueeze(0) - C_xy[0])).exp() * (α.unsqueeze(-1) @ β.unsqueeze(0)))[0]
p_xx, p_yy = None, None
if debias:
  p_xx = (((1 / ε) * (f_aa.unsqueeze(-1) + f_aa.unsqueeze(0) - C_xx[0])).exp() * (α.unsqueeze(-1) @ α.unsqueeze(0)))[0]
  p_yy = (((1 / ε) * (g_bb.unsqueeze(-1) + g_bb.unsqueeze(0) - C_yy[0])).exp() * (β.unsqueeze(-1) @ β.unsqueeze(0)))[0]

Thanks, Arnaud

[jeanfeydy]

Hi @arnaudvl,

Thanks :-) I fully agree with you, and have started to work on a much better API for GeomLoss in the ot_api branch. I will focus on this in the last two weeks of May: hopefully, I will finally be able to make a new release for GeomLoss in early June that solves this issue!

Best regards, Jean

[arnaudvl]

Great, the examples you linked look like a very user friendly way indeed. Looking forward to the release!

[ursueugen]

Hi @jeanfeydy and thanks for the great package!

Any plans to implement the return of the optimal transport plan?

Best regards, Eugen