ott icon indicating copy to clipboard operation
ott copied to clipboard

restructuring neural models + addition of OT-FM and GENOT

Open MUCDK opened this issue 1 year ago • 10 comments

This is a PR for

  1. new base classes neural solvers and models (i.e. neural networks)
  2. Incorporating unbalancedness and learning the rescaling factors for any neural OT model.
  3. adding OTFlowMatching and, related to this, classes for flows and time samplers
  4. adding GENOT (with extension to conditional GENOT)
  5. adding drafts of data loaders.

Following this PR, the implementations of ICNN-based solvers and the Monge Gap model should be adapted and extended to the unbalanced setting.

~~Moreover, wrt typing, I replaced jnp.ndarray by jax.Array~~

What remains to be done, but I would prefer to do in a separate PR

  1. add graph costs to OTFM and GENOT, i.e. functions which compute batch-wise graphs, and compute costs, e.g. geodesic Sinkhorn or convolutional Wasserstein from this.
  2. implementations of ICNN-based solvers and the Monge Gap model should be adapted and extended to the unbalanced setting.

MUCDK avatar Nov 22 '23 09:11 MUCDK