Formatting fixes using Black

Adds tests for overlapping context bonds, and a case with three modalities. context_bonds and dim_cond are now pytest.parametrize arguments instead of hardcoded class attributes.

Author: DhruvaRajwade

cmonge/models/nn.py

@@ -324,7 +324,11 @@ class ConditionalPerturbationNetwork(BasePotential):
 
     @nn.compact
     def __call__(
-        self, x: jnp.ndarray, c: jnp.ndarray, num_contexts: int = 2, deterministic: bool = True
+        self,
+        x: jnp.ndarray,
+        c: jnp.ndarray,
+        num_contexts: int = 2,
+        deterministic: bool = True,
     ) -> jnp.ndarray:  # noqa: D102
         """
         Args:
@@ -387,25 +391,37 @@ def __call__(
                 stacked_embeddings = jnp.stack(embeddings, axis=1)  # (Batch, N, Dim)
 
                 # Input Dropout
-                stacked_embeddings = nn.Dropout(rate=self.dropout_rate, deterministic=deterministic)(stacked_embeddings)
+                stacked_embeddings = nn.Dropout(
+                    rate=self.dropout_rate, deterministic=deterministic
+                )(stacked_embeddings)
 
                 # Multi-Head Attention Scores
-                att_layer = nn.Dense(self.num_heads, use_bias=True, name="AttentionScores")
+                att_layer = nn.Dense(
+                    self.num_heads, use_bias=True, name="AttentionScores"
+                )
                 scores = att_layer(stacked_embeddings)  # (Batch, N, Heads)
                 weights = jax.nn.softmax(scores, axis=1)
 
                 # Attention Weights Dropout
-                weights = nn.Dropout(rate=self.dropout_rate, deterministic=deterministic)(weights)
+                weights = nn.Dropout(
+                    rate=self.dropout_rate, deterministic=deterministic
+                )(weights)
 
                 # Weighted Pooling: (B, N, D), (B, N, H) -> (B, H, D)
-                weighted_sum = jnp.einsum('bnd,bnh->bhd', stacked_embeddings, weights)
+                weighted_sum = jnp.einsum("bnd,bnh->bhd", stacked_embeddings, weights)
 
                 # Flatten and Project
-                cond_embedding = weighted_sum.reshape(weighted_sum.shape[0], -1)  # (B, H*D)
-                cond_embedding = nn.Dense(dim_cond_map[0], use_bias=True, name="AttentionOutput")(cond_embedding)
+                cond_embedding = weighted_sum.reshape(
+                    weighted_sum.shape[0], -1
+                )  # (B, H*D)
+                cond_embedding = nn.Dense(
+                    dim_cond_map[0], use_bias=True, name="AttentionOutput"
+                )(cond_embedding)
 
                 # Output Dropout
-                cond_embedding = nn.Dropout(rate=self.dropout_rate, deterministic=deterministic)(cond_embedding)
+                cond_embedding = nn.Dropout(
+                    rate=self.dropout_rate, deterministic=deterministic
+                )(cond_embedding)
             else:
                 # Average along stacked dimension (alternatives like summing are possible)
                 cond_embedding = jnp.mean(jnp.stack(embeddings), axis=0)
@@ -434,7 +450,7 @@ def create_train_state(
 
         # Split rng for dropout keys during init
         rng, rng_dropout = jax.random.split(rng)
-        init_rngs = {'params': rng, 'dropout': rng_dropout}
+        init_rngs = {"params": rng, "dropout": rng_dropout}
 
         params = self.init(init_rngs, x=x, c=c)["params"]
         return PotentialTrainState.create(

cmonge/tests/models/test_attention_pooling.py

@@ -1,104 +1,148 @@
+import pytest
 import jax
 import jax.numpy as jnp
-import optax
 
 from cmonge.models.nn import ConditionalPerturbationNetwork
 
+# (context_bonds, dim_cond, num_contexts)
+CONTEXT_BOND_CONFIGS = [
+    pytest.param(
+        ((0, 10), (10, 20)),
+        20,
+        2,
+        id="non_overlapping_2_modalities",
+    ),
+    pytest.param(
+        ((0, 10), (0, 10)),
+        10,
+        2,
+        id="overlapping_2_modalities",
+    ),
+    pytest.param(
+        ((0, 10), (10, 20), (20, 30)),
+        30,
+        3,
+        id="non_overlapping_3_modalities",
+    ),
+]
+
+DIM_DATA = 16
+DIM_HIDDEN = [32, 32]
+DIM_COND_MAP = (8,)
+BATCH_SIZE = 4
+
+
+def _make_model(context_bonds, attention_pooling, dropout_rate=0.1):
+    return ConditionalPerturbationNetwork(
+        dim_hidden=DIM_HIDDEN,
+        dim_data=DIM_DATA,
+        dim_cond=max(stop for _, stop in context_bonds),
+        dim_cond_map=DIM_COND_MAP,
+        embed_cond_equal=True,
+        attention_pooling=attention_pooling,
+        num_heads=4,
+        dropout_rate=dropout_rate,
+        context_entity_bonds=context_bonds,
+    )
+
+
+def _make_inputs(rng, dim_cond):
+    rng_x, rng_c = jax.random.split(rng)
+    x = jax.random.normal(rng_x, (BATCH_SIZE, DIM_DATA))
+    c = jax.random.normal(rng_c, (BATCH_SIZE, dim_cond))
+    return x, c
+
 
 class TestAttentionPooling:
     """Tests for attention pooling in ConditionalPerturbationNetwork."""
 
-    # Shared config for a model using embed_cond_equal (deep set path)
-    DIM_DATA = 16
-    DIM_COND = 20  # 2 contexts of size 10 each
-    DIM_HIDDEN = [32, 32]
-    DIM_COND_MAP = (8,)
-    CONTEXT_BONDS = ((0, 10), (10, 20))
-    BATCH_SIZE = 4
-    NUM_CONTEXTS = 2
-
-    def _make_model(self, attention_pooling: bool, dropout_rate: float = 0.1):
-        return ConditionalPerturbationNetwork(
-            dim_hidden=self.DIM_HIDDEN,
-            dim_data=self.DIM_DATA,
-            dim_cond=self.DIM_COND,
-            dim_cond_map=self.DIM_COND_MAP,
-            embed_cond_equal=True,
-            attention_pooling=attention_pooling,
-            num_heads=4,
-            dropout_rate=dropout_rate,
-            context_entity_bonds=self.CONTEXT_BONDS,
-        )
-
-    def _make_inputs(self, rng):
-        rng_x, rng_c = jax.random.split(rng)
-        x = jax.random.normal(rng_x, (self.BATCH_SIZE, self.DIM_DATA))
-        c = jax.random.normal(rng_c, (self.BATCH_SIZE, self.DIM_COND))
-        return x, c
-
-    def test_attention_pooling_forward_pass(self):
+    @pytest.mark.parametrize(
+        "context_bonds,dim_cond,num_contexts", CONTEXT_BOND_CONFIGS
+    )
+    def test_attention_pooling_forward_pass(
+        self, context_bonds, dim_cond, num_contexts
+    ):
         """Test that attention pooling produces correct output shape."""
-        model = self._make_model(attention_pooling=True)
+        model = _make_model(context_bonds, attention_pooling=True)
         rng = jax.random.PRNGKey(0)
-        x, c = self._make_inputs(rng)
+        x, c = _make_inputs(rng, dim_cond)
 
         rng_params, rng_dropout = jax.random.split(rng)
-        params = model.init(
-            {"params": rng_params, "dropout": rng_dropout}, x=x, c=c
-        )["params"]
+        params = model.init({"params": rng_params, "dropout": rng_dropout}, x=x, c=c)[
+            "params"
+        ]
 
-        out = model.apply({"params": params}, x, c, self.NUM_CONTEXTS)
-        assert out.shape == (self.BATCH_SIZE, self.DIM_DATA)
-        # Output should be a residual: x + f(x, c), so not all zeros
+        out = model.apply({"params": params}, x, c, num_contexts)
+        assert out.shape == (BATCH_SIZE, DIM_DATA)
         assert not jnp.allclose(out, 0.0)
 
-    def test_both_pooling_modes_same_output_shape(self):
-        """Test that mean pooling and attention pooling produce the same output shape."""
+    @pytest.mark.parametrize(
+        "context_bonds,dim_cond,num_contexts", CONTEXT_BOND_CONFIGS
+    )
+    def test_both_pooling_modes_same_output_shape(
+        self, context_bonds, dim_cond, num_contexts
+    ):
+        """Test that mean and attention pooling produce the same output shape."""
         rng = jax.random.PRNGKey(42)
-        x, c = self._make_inputs(rng)
+        x, c = _make_inputs(rng, dim_cond)
 
-        # Mean pooling (default)
-        model_mean = self._make_model(attention_pooling=False)
+        model_mean = _make_model(context_bonds, attention_pooling=False)
         rng_p1, rng_d1, rng_p2, rng_d2 = jax.random.split(rng, 4)
-        params_mean = model_mean.init(
-            {"params": rng_p1, "dropout": rng_d1}, x=x, c=c
-        )["params"]
-        out_mean = model_mean.apply({"params": params_mean}, x, c, self.NUM_CONTEXTS)
-
-        # Attention pooling
-        model_attn = self._make_model(attention_pooling=True)
-        params_attn = model_attn.init(
-            {"params": rng_p2, "dropout": rng_d2}, x=x, c=c
-        )["params"]
-        out_attn = model_attn.apply({"params": params_attn}, x, c, self.NUM_CONTEXTS)
-
-        assert out_mean.shape == out_attn.shape == (self.BATCH_SIZE, self.DIM_DATA)
-
-    def test_dropout_deterministic_vs_stochastic(self):
-        """Test that deterministic=False (training) produces different outputs across runs
-        while deterministic=True (eval) is consistent."""
-        model = self._make_model(attention_pooling=True, dropout_rate=0.5)
+        params_mean = model_mean.init({"params": rng_p1, "dropout": rng_d1}, x=x, c=c)[
+            "params"
+        ]
+        out_mean = model_mean.apply({"params": params_mean}, x, c, num_contexts)
+
+        model_attn = _make_model(context_bonds, attention_pooling=True)
+        params_attn = model_attn.init({"params": rng_p2, "dropout": rng_d2}, x=x, c=c)[
+            "params"
+        ]
+        out_attn = model_attn.apply({"params": params_attn}, x, c, num_contexts)
+
+        assert out_mean.shape == out_attn.shape == (BATCH_SIZE, DIM_DATA)
+
+    @pytest.mark.parametrize(
+        "context_bonds,dim_cond,num_contexts", CONTEXT_BOND_CONFIGS
+    )
+    def test_dropout_deterministic_vs_stochastic(
+        self, context_bonds, dim_cond, num_contexts
+    ):
+        """Test that deterministic=False produces different outputs across runs
+        while deterministic=True is consistent."""
+        model = _make_model(context_bonds, attention_pooling=True, dropout_rate=0.5)
         rng = jax.random.PRNGKey(7)
-        x, c = self._make_inputs(rng)
+        x, c = _make_inputs(rng, dim_cond)
 
         rng_params, rng_dropout = jax.random.split(rng)
-        params = model.init(
-            {"params": rng_params, "dropout": rng_dropout}, x=x, c=c
-        )["params"]
+        params = model.init({"params": rng_params, "dropout": rng_dropout}, x=x, c=c)[
+            "params"
+        ]
 
         # Deterministic mode: two calls should be identical
-        out_eval_1 = model.apply({"params": params}, x, c, self.NUM_CONTEXTS, deterministic=True)
-        out_eval_2 = model.apply({"params": params}, x, c, self.NUM_CONTEXTS, deterministic=True)
+        out_eval_1 = model.apply(
+            {"params": params}, x, c, num_contexts, deterministic=True
+        )
+        out_eval_2 = model.apply(
+            {"params": params}, x, c, num_contexts, deterministic=True
+        )
         assert jnp.allclose(out_eval_1, out_eval_2)
 
         # Stochastic mode: two calls with different dropout keys should differ
         key1, key2 = jax.random.split(jax.random.PRNGKey(99))
         out_train_1 = model.apply(
-            {"params": params}, x, c, self.NUM_CONTEXTS,
-            deterministic=False, rngs={"dropout": key1},
+            {"params": params},
+            x,
+            c,
+            num_contexts,
+            deterministic=False,
+            rngs={"dropout": key1},
         )
         out_train_2 = model.apply(
-            {"params": params}, x, c, self.NUM_CONTEXTS,
-            deterministic=False, rngs={"dropout": key2},
+            {"params": params},
+            x,
+            c,
+            num_contexts,
+            deterministic=False,
+            rngs={"dropout": key2},
         )
         assert not jnp.allclose(out_train_1, out_train_2)

cmonge/trainers/conditional_monge_trainer.py

@@ -192,7 +192,7 @@ def loss_fn(
                 batch["source"],
                 batch["condition"],
                 n_contexts,
-                **kwargs
+                **kwargs,
             )
 
             # compute the loss