Self-edges are never zeroed in MaskedMessagePassingLayer (unassigned .at[].set() is a no-op)

[AMX07]

In src/oneqmc/wf/nn/masked/message_passing.py (lines 60–63):

# We will set self-edges to 0
edge_part.at[
    ..., jnp.arange(edge_part.shape[-4]), jnp.arange(edge_part.shape[-3]), :, :
].set(0.0)

JAX arrays are immutable: .at[].set() returns a new array, and here the result is discarded, so this statement is a silent no-op. Contrary to the comment, self-edges (i = j messages) are never zeroed and contribute to every message-passing update.

Minimal repro:

import jax.numpy as jnp

edge_part = jnp.ones((3, 3, 2, 4))
edge_part.at[..., jnp.arange(3), jnp.arange(3), :, :].set(0.0)
print(edge_part[0, 0].sum())   # 8.0 — unchanged

This affects MaskedMessagePassingLayer wherever it is used (the NucleiGNN and the OrbitalGenerator in Orbformer). Note that the self-edge features are not zero: the edge featurization of a zero diff vector still produces nonzero scalar features (exp(-r) = 1, sigmoid(s - r)/s), so live self-edges genuinely change the output.

Important caveat for any fix: the published LAC checkpoint was trained with this behavior, so simply adding the assignment (edge_part = edge_part.at[...].set(0.0)) would change the effective architecture and invalidate the released checkpoint. I discovered this while porting Orbformer to Apple MLX, where reproducing the checkpoint's per-geometry parameters to ~1e-6 against fp64 JAX outputs (across five geometries) required keeping self-edges live — so the shipped weights demonstrably have this baked in. The options seem to be:

1. delete the dead statement and fix the comment, documenting that self-edges are (by accident, now by contract) part of the trained model; or
2. apply the real fix behind a flag that defaults to the current behavior, so existing checkpoints keep working and only future training runs zero self-edges.

[ae-foster]

Great catch Ansh MIttal (@AMX07) . You're right that we would have to re-pretrain a separate checkpoint with this fixed for everything to work as intended