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

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.

Start in src/oneqmc/wf/nn/masked/message_passing.py at lines 60–63 and inspect the MaskedMessagePassingLayer call sites in NucleiGNN and OrbitalGenerator. Verify the discarded JAX update and compare the behavior required by the published LAC checkpoint. Done means the selected compatibility behavior is implemented or documented consistently for existing and future models.