Hi, this is a fantastic addition! I am currently trying to use sinkhorn_unbalanced2 and would love to have mass=True turned on for the marginal penalisation, but I am struggling to get it to work. Any tips on how to achieve it without much hacking?

(Sorry if I am being confused here -- potentially this does not affect the optimisation routine to find the optimal plan and so it really is just about returning the correct UOT cost?)