refactor(pallas): the fused complex-M2L helpers, where a broadcast hides a bug

[TobiBu]

Third of the three modules in this pass; see the re-measurement note at the bottom, which
changed what this PR is.

The mechanism: three helpers that reduce by broadcast

This module carries real and imaginary parts as separate real arrays rather than as a
complex dtype, so every reduction is written twice and the two halves have to agree by
hand. All three helpers reduce with an explicit broadcast rather than a matmul:

_matvec(mat, vec)     jnp.sum(mat * vec[None, :], axis=1)
_matvec_T(mat, vec)   jnp.sum(mat * vec[:, None],  axis=0)
_block_matmul(...)    jnp.sum(block_r * vec_r[:, None, :], -1) - ...(block_i)

A broadcast accepts a length-1 operand and spreads it. Measured on main:

input	main returns
length-1 vec into a (32, 16) operator	a full (32,) result
_matvec_T with a flattened (512,) operator	(512,)
vec_i of (4, 1) beside vec_r of (4, 8)	a plausible (4, 8), first column's imaginary part in every lane

The third is the dangerous one — a real/imag mismatch inside a complex kernel — and the
pilot never tried it: its perturbations are leading/trailing/extra/flattened, never
length-1. So it is a gap this closes beyond what was measured.

Axes by execution, read per CALL

Two sets of extents cannot establish a relation, so the 2026-09-07 recording was read
call by call:

_matvec    mat (16, 32) vec (32,)  |  mat (32, 16) vec (16,)
           -> `vec` is mat's SECOND axis, at two extents with the roles SWAPPED
_matvec_T  mat (16, 32) vec (16,)  |  mat (32, 16) vec (32,)
           mat (32, 128) vec (32,) |  mat (128, 32) vec (128,)
           -> `vec` is mat's FIRST axis, at four extents

rows/cols therefore name a relation between two arguments, not a width — one helper
is applied to three different operators at four extents — and the two new vocabulary
entries say exactly that.

What this PR does NOT claim

Worth reading before the diff, because two plausible claims are false:

• Swapping _matvec/_matvec_T was already caught. A (32,) vector will not broadcast
  against a 16-long axis, so main already raises — as TypeError rather than
  TypeCheckError, which is the only thing these annotations change about it. That test
  accepts either exception. Asserting TypeCheckError alone would dress an exception-type
  change up as a closed gap, and it did exactly that in my first draft.
• Two of the pilot's six acceptances on this pair are not defects. They perturbed rows
  and cols where nothing else binds them, and a (31, 16) operator with a (16,) vector is a
  perfectly well-formed matvec. A test asserts they stay accepted, so the pilot's report
  cannot talk someone into cross-binding an axis the evidence does not support.
• _block_matmul_vjp is deliberately left bare. It rejected all six of its
  perturbations, so §4.1 says leave it alone. Its recording is still what supplies the
  second extent for squareness — (4, 8, 8) and (8, 16, 16) — which _block_matmul's own
  single recorded extent cannot.

The Pallas lane still traces

_m2l_one reaches these helpers from inside pallas_call, where the operands are
kernel tracers rather than plain arrays, so the decorator had to survive that. Interpret
mode agrees with the JAX twin to 5.6e-16. The beartype cost is paid at trace time, not
per element, since these are traced once under vmap/pallas_call.

Tests

tests/unit/pallas/test_fused_m2l_helper_axis_contracts.py — 8 tests, 5 red against
main. The 3 green ones are the control, the free-axis pin and the adjoint-name pin, and
each says in its own docstring that it is green on main and why.

Verification

pre-commit run --all-files                        all pass
JAX_ENABLE_X64=1 pytest                           1921 passed, 137 skipped
JAX_ENABLE_X64=1 pytest tests/characterization     34 passed  (golden unmoved)
JACCPOT_RUNTIME_TYPECHECK=1 pytest tests/unit/pallas
  tests/unit/operators/test_m2l_complex_fused_pallas.py ...   145 passed, 23 skipped

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

main    844 shaped / 1253 bare   40.2%      45 bare / 19 shaped  pallas/m2l_complex_fused.py
branch  852 shaped / 1245 bare   40.6%      37 bare / 27 shaped  pallas/m2l_complex_fused.py

Why this is a three-helper PR and not a module pass

The doc's rate for this module was 30%, measured 2026-09-04 — before 19f1539 landed.
Re-recorded on main 2026-09-07 it is 7% (122 perturbations, 9 accepted, coverage
8/0/2). Of those 9: 6 are the helpers above, 2 are the free axes described above, and 1 is
m2l_complex_fused_pallas's multipoles, which is not closable by annotation — sh is
bound by that parameter alone and the output's sh derives from it, so both move together
and stay consistent. _m2l_one and _m2l_one_vjp are UNREPLAYABLE (opaque t), hence
unmeasured; they are transitively protected anyway, since every reduction in them goes
through the helpers this PR annotates.

🤖 Generated with Claude Code