fix(pallas): check the fused leaf kernels' shapes against each other

[TobiBu]

The last of the 2026-09-07 pilot pass, and the one that is not an annotation PR. Follows
#336 and #337; closes the nine acceptances those two deliberately left alone.

Why a guard and not an annotation

Annotating the position parameters would have closed all nine. It was still the wrong
instrument, and this is the harder version of checklist item 13 rather than a
straightforward application of it:

• each of these parameters' bodies raises a documented ValueError over its shape, and
  a decorator runs before the body, so the annotation replaces it;
• DELIBERATELY_BARE's first entry states the policy verbatim — changing which exception a
  caller sees is a behaviour change, not a docs change;
• but unlike refactor(runtime): finish item 5 -- geometry, node-scalar and per-order families #310's case, the annotation here would have closed something real, which is
  what made this a decision rather than a rule lookup.

So the body now checks what its own docstring already promised, under the exception callers
already expect. No annotations are added. The census moves by exactly +1 bare — the new
helper's array: Array parameter, which by design accepts any array and so genuinely cannot
pin a shape. Stated rather than gamed.

What the old checks missed

Every entry point tested ndim != 3 or shape[-1] != 3 and never the leading extent,
while the docstrings promised ValueError "if the input shapes are mutually inconsistent"
— a consistency the bodies did not verify. The docstring over-promised, and the gap it left
splits by severity:

perturbation	what happens	visible?
a position array one leaf short	the Pallas grid shrinks with it, so the output shrinks too	yes, to a caller who checks
a mask one row or column short	the grid is unchanged, the BlockSpec indexes the mask out of bounds, JAX clamps, and leaf 3 reuses leaf 2's mask	no — right shape, wrong numbers

The second is the one worth the PR. Measured on main: target_mask one row short returned
a full (3, 2, 4); so did one column short; so did source_positions one leaf short.

The annotations could not have caught the mask cases anyway without the positions:
target_mask: Bool[Array, "leaves w"] beside a bare target_positions binds leaves
and w with one parameter, so those names have nothing to disagree with. The source side
was already safe precisely because source_masses and source_mask both carry
srcleaves srcslots and cross-check each other — the pilot rejected those, and this PR does
not claim them.

The reference lane needed one too, for a different reason

nearfield_leafpair_jax had no shape check at all. A leaf_positions of trailing width
2 makes accels two-wide, and the final concatenate with the potential returns a
three-wide result — the rank is right, the contract is 4, and nothing raised. (3, 2, 2, 3)
in, (3, 2, 3) out. This is the lane the Pallas kernels are checked against, so it
mattered more than its one acceptance suggests.

Ordering is part of the check, and I got it wrong first

Worth reading, because it is the same mistake as item 13 committed in the other direction.
Placing the generic source-table checks before #297's source-width guard made that
guard's specific message unreachable — it explains that a narrower pool reads out of bounds
while a wider one silently drops real particles, which a generic message cannot — and turned
five existing tests red:

test_the_decoupled_source_pool_is_its_own_leading_axis
test_decoupled_rejects_a_source_pool_of_a_different_width[narrower-8-4]
test_decoupled_rejects_a_source_pool_of_a_different_width[narrower-16-3]
test_decoupled_rejects_a_source_pool_of_a_different_width[wider-4-8]
test_decoupled_rejects_a_source_pool_of_a_different_width[wider-4-16]

They were fixed by reordering the checks, not by touching the tests. A new test pins the
priority so a future reorder is visible.

The decoupled variant's independence is preserved

num_targets and num_sources are deliberately independent there — that separation is the
whole point of the lane — so each operand is checked against the array it belongs to rather
than against a single leaf count. Verified: a source pool of 2 leaves against 3 target leaves
still goes through, and the control test asserts it. A guard that tied them together would
have broken the lane while looking like an improvement.

Tests

tests/unit/pallas/test_fused_leaf_shape_guards.py — 11 tests, 9 red against main, one
per closed acceptance. The 2 green are the control (all four entry points, including the
decoupled independence) and the #297 priority pin.

The sibling test_fused_leaf_axis_contracts.py had a docstring asserting that "what the
bodies do not check is the masks, masses, ids and validity arrays". That is no longer true,
so it is updated; the two files now divide as annotations there, body guards here.

STYLE_GUIDE.md item 13 gains the general form: an annotation is the wrong instrument
wherever a parameter's own body documents a ValueError over its shape, even when it would
work — plus the ordering trap above.

Verification

pre-commit run --all-files                                     all pass
JAX_ENABLE_X64=1 pytest                                        1932 passed, 138 skipped
JAX_ENABLE_X64=1 pytest tests/characterization                  34 passed (golden unmoved)
JACCPOT_RUNTIME_TYPECHECK=1 pytest tests/unit/pallas
  tests/unit/operators tests/unit/nearfield
  tests/unit/test_custom_vjp_parity.py                          572 passed, 29 skipped
JACCPOT_RUNTIME_TYPECHECK=1 pytest tests/unit/test_type_annotation_guard.py   5 passed

Census — bench/annotation_census.py, quoted not retyped. Shaped count and share unchanged,
which is the point:

main    852 shaped / 1246 bare   40.6%      42 bare / 16 shaped  pallas/nearfield_fused_leaf.py
branch  852 shaped / 1247 bare   40.6%      43 bare / 16 shaped  pallas/nearfield_fused_leaf.py

🤖 Generated with Claude Code