[rssrwn]

Depends on #2211 (this PR is based on its branch, I'll retarget to main once it's merged).

This was a problem Claude noticed when I was added the unit tests for EdgeDegreeEmbedding. It is fairly niche, but has a significant impact on memory usage when activation checkpointing is turned on and running on large systems, so it seems worth fixing (the fix is also only 1 or 2 lines). This PR also adds a test based on the number of copies of the input tensor, this fails on the old code and passes with the fix in place. Below is the summary given by Claude.

Problem
With activation checkpointing, EdgeDegreeEmbedding.forward passes the full node tensor into every checkpointed edge chunk and gets a new copy back. Checkpointing keeps each chunk's inputs until backward, so one node-feature tensor is kept per chunk (131,072 edges each). Chunks grow with edges, so this memory grows roughly quadratically with system size. Edgewise avoids this by passing one shared tensor and summing per-chunk contributions.

Who hits it
Conservative force/stress inference or training on the general backend with activation checkpointing on, for systems with more than one chunk. Checkpointing is on with InferenceSettings() defaults (used by MLIPPredictUnit when no settings are passed), the batch preset, and configs that enable it.

Measured impact (one H100 80GB, uma-s-1p2, default InferenceSettings(), energy + forces on rattled Cu fcc; peak GPU memory of the call above the pre-call baseline)

atoms	edges	chunks	peak memory main / this PR	difference	time main / this PR
6,912	0.54M	5	11.05 / 10.92 GB	0.13 GB (1%)	1.34 / 1.34 s
19,652	1.53M	12	19.61 / 18.61 GB	1.00 GB (5%)	3.79 / 3.79 s
42,592	3.32M	26	37.38 / 32.46 GB	4.92 GB (15%)	8.24 / 8.25 s
78,732	6.14M	47	70.95 / 54.26 GB	16.69 GB (31%)	15.38 / 15.37 s

Each difference matches (chunks − 1) × one node-feature tensor. At 79k atoms main peaks at 72 GB on the first call. Extrapolating, main exceeds 80 GB at around 90k atoms for this system, while this PR fits up to roughly 125k.

Fix
Each chunk scatters onto a shared zero tensor and its contribution is added outside the checkpoint, so one node-sized tensor is kept, as in Edgewise.

Correctness and cost

• Forces match main to within 1.3e-6 ebove
• In fp64, outputs, forces and a Hessian-vector product agree with checkpointing off to ~1e-15 (18
  chunks)
• New regression test counts node-feature copies kept for backward (fails on main)
• No measurable time change on GPU; onbout 1% slower at the default chunksize, up to ~5.5% with many small chunks (one extra node-sized add per chunk)