fix(model): zero masked pair slots before the reverse gather

Masked pair slots left each GNN layer holding nonzero features, since the
transformer returns edge + attn + mlp(...) unmasked. The reverse gather is
the only consumer of a masked slot, so a reverse pointing at the padded
sentinel picked those features up, and the output then depended on the
packed width -- a caller-side shape choice, not physics.

Reachable wherever a reverse lands on the sentinel: under truncate_edges
overflow, or through pack_edges on a neighbour list trimmed asymmetrically
upstream (a per-center knn cap), where it happens with overflow=False and
no diagnostic.

Guard it structurally: padded slots must leave the backbone at exactly
zero, plus a width-invariance check over an orphaned reverse. Note that a
fresh init leaves every bias zero, which zeroes the sentinel and hides the
leak, so the invariance test shifts parameters off init.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: Egor Rumiantsev

src/petjax/model.py

@@ -216,7 +216,7 @@ def __call__(
                 )(node, edge, cutoffs_tokens, mask)
 
             # Message passing (feedforward mixing)
-            edge_flat = edge.reshape(P, d_pet)
+            edge_flat = edge.reshape(P, d_pet) * pair_mask[..., None]
             reversed_flat = edge_flat[reverse]
             combined = jnp.concatenate([edge_flat, reversed_flat], axis=-1)
             combined = nn.LayerNorm(name=f"comb_norms_{layer_idx}")(combined)

tests/test_select.py

@@ -112,3 +112,81 @@ def test_pack_edges_forward_smoke(structure):
     params = model.init(jax.random.key(0), **packed)
     energy = model.apply(params, **packed)
     assert bool(jnp.all(jnp.isfinite(energy)))
+
+
+def _orphaned(structure):
+    """Copy of `structure` with one direction of one pair dropped.
+
+    Mimics an upstream per-center trim (or an overflow drop): the surviving
+    direction keeps a `reverse` pointing at the padded sentinel pair, so the
+    forward gathers a masked slot.
+    """
+    structure = {k: (v.copy() if hasattr(v, "copy") else v) for k, v in structure.items()}
+    sentinel = structure["pair_mask"].shape[0] - 1
+    real = np.flatnonzero(structure["pair_mask"][:sentinel])
+    kept = int(real[0])
+    dropped = int(structure["reverse"][kept])
+    assert dropped != kept
+
+    structure["pair_mask"][dropped] = False
+    structure["reverse"][kept] = sentinel
+    return structure
+
+
+def _backbone():
+    from petjax.model import Backbone
+
+    return Backbone(
+        d_pet=8,
+        d_node=8,
+        d_feedforward=8,
+        num_heads=2,
+        num_attention_layers=1,
+        num_gnn_layers=2,
+        cutoff=CUTOFF,
+    )
+
+
+def _run_backbone(structure, k, shift=0.0):
+    """Backbone on `structure` packed to width `k`.
+
+    `shift` offsets every parameter: a fresh init leaves all biases zero, which
+    happens to zero the padded sentinel slot and hide any leak through it.
+    """
+    backbone = _backbone()
+    packed, overflow = _pack(structure, k)
+    assert not bool(overflow)
+    inputs = {key: value for key, value in packed.items() if key != "pair_cutoffs"}
+    params = backbone.init(jax.random.key(0), **inputs, pair_cutoffs=None)
+    if shift:
+        params = jax.tree.map(lambda x: x + shift, params)
+    node, messages, _ = backbone.apply(params, **inputs, pair_cutoffs=None)
+    return np.asarray(node), np.asarray(messages), np.asarray(packed["pair_mask"])
+
+
+def test_masked_slots_carry_no_features(structure):
+    """Padded pair slots leave the backbone at exactly zero.
+
+    The reverse gather is the one consumer of a masked slot, so anything left
+    there leaks into real edges whose reciprocal was trimmed upstream.
+    """
+    _, messages, pair_mask = _run_backbone(structure, _max_count(structure) + 5)
+    assert np.all(messages[~pair_mask] == 0.0)
+
+
+def test_backbone_invariant_to_packed_width(structure):
+    """`k` is a padding choice: widening it must not move real features.
+
+    Runs with an orphaned reverse -- the one case that reads a masked slot --
+    and with shifted parameters, so the sentinel slot is not incidentally zero.
+    """
+    orphaned = _orphaned(structure)
+    k = _max_count(orphaned)
+
+    node_tight, msg_tight, mask_tight = _run_backbone(orphaned, k, shift=0.3)
+    node_wide, msg_wide, mask_wide = _run_backbone(orphaned, k + 11, shift=0.3)
+
+    np.testing.assert_allclose(node_tight, node_wide, rtol=1e-5, atol=1e-6)
+    np.testing.assert_allclose(
+        msg_tight[mask_tight], msg_wide[mask_wide], rtol=1e-5, atol=1e-6
+    )