Small fixes checkpoint loading and ott-jax version bump (#31)

Co-authored-by: Alice Driessen <alicedriessen@Alices-MacBook-Air.local>
Co-authored-by: Jannis Born <jannis.born@gmx.de>

Author: 3 people

cmonge/trainers/ot_trainer.py

@@ -15,7 +15,7 @@
 from ott.neural.methods.monge_gap import MongeGapEstimator
 from ott.neural.methods.neuraldual import W2NeuralDual
 from ott.neural.networks.icnn import ICNN
-from ott.neural.networks.potentials import MLP
+from ott.neural.networks.potentials import PotentialMLP
 
 from cmonge.datasets.single_loader import AbstractDataModule
 from cmonge.evaluate import (
@@ -199,6 +199,7 @@ def setup(
         fitting_loss: Dict[str, Any],
         regularizer: Dict[str, Any],
         optim: Dict[str, Any],
+        checkpointing_path: Optional[str] = None,  # For compatibility with base class
     ) -> None:
         """Initializes models and optimizers."""
         self.metrics["params"] = {
@@ -218,7 +219,7 @@ def setup(
         regularizer = partial(regularizer_fn, **regularizer.kwargs)
 
         # setup neural network model
-        model = MLP(dim_hidden=dim_hidden, is_potential=False, act_fn=nn.gelu)
+        model = PotentialMLP(dim_hidden=dim_hidden, is_potential=False, act_fn=nn.gelu)
 
         # setup optimizer and scheduler
         opt_fn = optim_factory[optim.name]
@@ -313,7 +314,7 @@ def setup(
             dim_hidden=dim_hidden,
             gaussian_map_samples=(samples_source, samples_target),
         )
-        neural_g = MLP(dim_hidden=dim_hidden)
+        neural_g = PotentialMLP(dim_hidden=dim_hidden)
 
         lr_schedule = optax.cosine_decay_schedule(
             init_value=lr, decay_steps=num_train_iters, alpha=1e-2

cmonge/utils.py

@@ -20,6 +20,7 @@ def monge_get_source_target_transport(
     source=True,
     transport=True,
     batch_size: int = None,
+    num_contexts: int = 2,
 ):
     if batch_size is None:
         batch_size = datamodule.batch_size
@@ -86,7 +87,10 @@ def monge_get_source_target_transport(
             )
 
         if transport:
-            trans = trainer.transport(source_expr, num_contexts=2)
+            if num_contexts == 0:
+                trans = trainer.transport(source_expr)
+            else:
+                trans = trainer.transport(source_expr, num_contexts=num_contexts)
             trans = datamodule.decoder(trans)
             trans_meta = cond_meta.copy()
             trans_meta["dtype"] = "transport"