Adding conditional monge to ott-jax

[DriessenA]

Dear developers,

Problem statement
Current Monge Gap implementation doesn't extend to unseen/novel conditions

Solution
We've implemented a conditional monge gap as extension of the monge gap from ott-jax, see paper and code. This implementation can be applied to unseen conditions (and cells, like the already implemented Monge Gap). We would love to add this to the ott-jax library if there is interest in this.

[DriessenA]

@jannisborn

[michalk8]

Hi @DriessenA , yes, that would amazing, thanks a lot! If you will open a PR, will be happy to review it!

[DriessenA]

Hi Michael, currently working on the PR. There is currently no standard dataset/dataloader that I should use for the conditional monge gap, correct? I see the OTDataset, but it's not used (as far as I could find), nor is it a dataloader.