refactor(m2l): the chunked scatter's shared axes -- and the one that is free

[TobiBu]

Second of two PRs from the 2026-09-07 re-recording, and the one that closes
runtime/kernels/_m2l.py out. Companion to #336.

The module is at 1%, and that is the headline

The phase-2 doc had this module at 8% of 286 perturbations and advised skipping it.
Re-recorded on main after cbc99f1: 244 perturbations, 2 accepted — 1%, on 15
measured functions with zero inconclusive and zero unreplayable. The class and accumulator
families are closed, and both survivors are in one function.

Only one of the two is a defect. Separating them is most of this change.

The defect: a gather that JAX clamps

_chunk_segment_scatter_add does

sort_idx        = jnp.argsort(masked_targets)   # length comes from tgt_chunk
contribs_sorted = contribs[sort_idx]            # <-- gathered with THAT length

so a contribs one row short is an out-of-bounds gather, and JAX clamps rather than
raising. Measured on main: the call returned a full (255, 25) accumulator having
silently gathered row 510 twice, scattering one pair's M2L contribution into the wrong
target. chunkflat now ties contribs, tgt_chunk and valid together, and closes it.

The non-defect: a free axis, proven not argued

local_accum's leading axis. 12 recorded calls show it at 7, 15, 31, 127, 255, 511 and
1023 against an unchanged contribs — it is the target-node count, which has nothing to
do with the chunk width. The pilot was perturbing an axis bound by one parameter with
nothing deriving from it, which produces a well-formed call.

So test_the_accumulators_node_axis_stays_free asserts it stays accepted. Without that,
the pilot's own report is a standing invitation to cross-bind an axis the evidence
contradicts.

Evidence for the two axes that are real

Read per call, 12 recordings, 9 distinct combinations:

axis	evidence
chunkflat	contribs leading == tgt_chunk == valid in all 9, at two extents (512, 4096)
sh	contribs trailing == local_accum trailing in all 9, at four extents (4, 9, 25, 81)

sh was already rejected on main by broadcasting, so its test accepts either exception
and its docstring says so: the annotation pins the name, it does not close a hole.

Inexact, not Float

The recording carries complex128, complex64 and float64 through both contribs
and local_accum. Narrowing to Float would reject the complex basis outright — the
mistake #293 shipped one module over, from a real-basis-only recording, for 27 CI failures.
A parametrised control test now covers both bases.

The doc commit

The second commit rewrites the phase-2 document's rows for all three modules in this pass
(30/29/8% → 7/8/1%) and names two categories it did not have:

• the Pallas clamp — 9 of the pass's 20 acceptances (a third commit corrects this: I
  first wrote 17, which was 20 minus the 3 free axes, a different quantity). A pallas_call
  takes its grid from one array's shape and indexes every other operand with the same block
  index, so a short operand is clamped rather than rejected. The pure-JAX twins reject the identical
  perturbation. nearfield_fused_leaf_jax accepted 0 of 21 where its _pallas twin
  accepted 4, which makes the Pallas lane strictly weaker at shape validation than the lane
  _m2l.py's docstring requires it to equal;
• the free axis — the category above, 3 of the 20, so future passes stop counting them
  as backlog.

It also records the prediction tally, 4 right / 4 wrong / 1 partial, because the misses
share a direction: I predicted acceptance from "nothing annotates it" and forgot that
ordinary array arithmetic rejects most shape errors by failing to broadcast. The four
correct predictions were all about broadcast reductions, where it does not.

Tests

Appended to tests/unit/runtime/test_m2l_shape_contracts.py — 5 tests, 1 red against
main. That ratio is the honest one for a module at 1%; the other four are two controls and
two deliberate pins, each labelled as green-on-main in its own docstring.

Verification

pre-commit run --all-files                        all pass
JAX_ENABLE_X64=1 pytest                           1918 passed, 137 skipped
JAX_ENABLE_X64=1 pytest tests/characterization     34 passed  (golden unmoved)
JACCPOT_RUNTIME_TYPECHECK=1 pytest tests/unit/test_type_annotation_guard.py   5 passed

Census — bench/annotation_census.py, quoted not retyped:

main    844 shaped / 1253 bare   40.2%      65 bare / 27 shaped  runtime/kernels/_m2l.py
branch  848 shaped / 1249 bare   40.4%      61 bare / 31 shaped  runtime/kernels/_m2l.py

Note on the .pre-commit-config.yaml conflict with #336

Both PRs add to the flake8 --builtins list: this one chunkflat, #336 rows,cols.
Whichever merges second wants the union of the two, which is the same resolution used in
this programme twice already.

🤖 Generated with Claude Code