perf(reg_split): compact the ConstantSplit split-stencil scatter (DelaunayNN 12x on A100)

[Jammy2211]

Summary

The JAX path of pixel_splitted_regularization_matrix_from scattered the full
(4P, K, K) outer product of the split stencil. For the natural-neighbour
DelaunayNN mesh the stencil tables are padded to K = 33
(SIBSON_MAX_NEIGHBORS 32 + 1 self column) while the real occupied width on a
production HST cell is min 1 / median 5 / p99 9 / max 11 — so ~97% of the 6.5M
scattered entries were padding, at a cost quadratic in the padded width. On the
A100 this single scatter was the largest remaining block of the DelaunayNN
params -> H prefix.

This PR scatters only the first kc = min(K, 12) columns and supplements the
256 widest rows with the blocks the compact pass did not cover. The result is
exact, not approximate: padded columns carry mapping 0 / weight 0, so a
row whose occupied size is <= kc is reproduced bit-for-bit by the compact pass
alone, and the wide rows get the head x tail, tail x head and tail x tail blocks
back in one masked .at[].add. K <= 12 callers (the Delaunay mesh's K = 4,
and the adapt-split family) take today's single scatter unchanged, with no
supplement and no guard.

Design

• Compact main scatter, width 12. SPLIT_REG_COMPACT_WIDTH = 12. Width was
  chosen from the measured cost curve, not guessed — see the width sweep below.
• Top-256 widest-row supplement. SPLIT_REG_WIDE_ROW_BUDGET = 256. Rows are
  selected by jax.lax.top_k on the integer splitted_sizes, which is not
  differentiated through, so the weights stay fully differentiable. The
  supplement costs W * (K**2 - kc**2) ~ 0.24M entries, an order of magnitude
  below the 6.5M it replaces.
• NaN on budget overflow. If more rows exceed kc than the budget holds, the
  matrix is poisoned with NaN rather than silently truncated. This is the same
  convention the Sibson natural-neighbour caps already use
  (mesh/interpolator/sibson.py: neighbor_overflow / failed set the
  interpolation weights to NaN) — the likelihood evaluates to NaN and the sample
  is discarded by the sampler, never returning a silently wrong H.
• The budget is sized against a real ensemble. The cap audit over 101
  ensemble geometries (4800 split rows each) sees a max of 50 rows wider than 12
  and a mean of 2.8 — a ~5x margin to the 256 budget. The widest single stencil
  in that ensemble is 21 natural neighbours (99.9th pct 11, 99.99th pct 15).
• NumPy path untouched.

A100 A/B (jobs 342334-342338, same node, same window)

measurement	before	after	speedup
assembly row (ms/call)	10.03	0.84	12.0x
params -> H prefix, vmap 16 (ms/call)	16.423	7.260	2.26x
params -> H prefix, unbatched (ms)	24.341	15.189	1.60x
full likelihood, single JIT (ms)	75.6	66.1	1.14x
full likelihood, vmap 16 (ms/call)	50.04	40.86	1.22x

Numerics: the pinned likelihood 29144.581944 is bit-identical on both legs
and pinned_drift: []. The pre-registered witness (assembly < 3 ms,
params -> H < 11 ms) is met with margin.

Why width 12, and why not the alternatives

Investigation jobs 342331/342332 priced every candidate on the same real HST
tables before the design was fixed:

approach	A100 fp64 ms/call	verdict
compact scatter, width 12 (this PR)	0.58	chosen
dense GEMM assembly	10.1	a wash with today's scatter
BCOO sparse assembly	21	2x worse
dedup / segment-sum assembly	39	4x worse
today's scatter (width 33)	10.03	baseline

Compact-width sweep (vmap 16, real HST tables):

width	12	16	20	24	28	32	33 (today)
A100 fp64 ms/call	0.58	1.47	2.78	4.53	6.71	9.31	10.03
CPU fp64 ms/call	12.6						72.8

Width 12 is 17x on GPU and 5.8x on CPU over the uncompacted scatter, so no
backend gate is needed.

API Changes

No breaking changes. pixel_splitted_regularization_matrix_from gains two
optional keyword arguments, compact_width and wide_row_budget, defaulting to
the new module constants SPLIT_REG_COMPACT_WIDTH = 12 and
SPLIT_REG_WIDE_ROW_BUDGET = 256; both are ignored on the NumPy path. Every
existing call site is unchanged.

One new behaviour on the JAX path: a geometry with more than 256 split rows
wider than 12 columns now yields a NaN regularization matrix (hence a NaN
likelihood, discarded by the sampler) instead of a silently narrower one. This
is the existing Sibson cap convention and is never reached on the audited
ensemble (max 50 wide rows of a 256 budget). The NumPy path is unchanged.

See full details below.

Test Plan

• pytest test_autoarray/inversion -q — 507 passed
• full pytest test_autoarray — 1468 passed
• New test_autoarray/inversion/regularizations/test_pixel_splitted_jax.py
  (8 tests): compact-vs-full parity on padded tables, K <= kc no-op path,
  wide-row supplement exactness, NaN-on-overflow, NumPy/JAX agreement.
• autolens_workspace_test scripts/misc/jax_assertions/delaunay_nn.py
  compaction parity section — PASS.
• Cap audit scripts/misc/jax_assertions/delaunay_nn_caps.py over 101
  ensemble geometries: max 50 / mean 2.8 wide rows of 4800, budget 256.
• A100 A/B, pinned likelihood bit-identical, pinned_drift: [].
• scripts/imaging/jax_grad/delaunay.py and
  scripts/imaging/jax_likelihood/delaunay.py (both reg.AdaptSplit on the
  JAX path, i.e. the changed function) run clean on this branch.

Links

• Results note: autolens_profiling/results/notes/delaunay_nn_constant_split_assembly.md
• Refs perf: compact the ConstantSplit split-stencil scatter (DelaunayNN) #536

Full API Changes (for automation & release notes)

Added

• autoarray.inversion.regularization.regularization_util.SPLIT_REG_COMPACT_WIDTH — module constant, 12; the number of stencil columns scattered for every row on the JAX path.
• autoarray.inversion.regularization.regularization_util.SPLIT_REG_WIDE_ROW_BUDGET — module constant, 256; the number of widest rows given a full-width supplementary scatter on the JAX path.

Changed Signature

• regularization_util.pixel_splitted_regularization_matrix_from(regularization_weights, splitted_mappings, splitted_sizes, splitted_weights, xp=np, compact_width=SPLIT_REG_COMPACT_WIDTH, wide_row_budget=SPLIT_REG_WIDE_ROW_BUDGET) — two new optional keyword arguments; no existing argument added, removed, renamed or re-defaulted.

Changed Behaviour

• regularization_util.pixel_splitted_regularization_matrix_from (JAX path only) — the (4P, K, K) scatter is replaced by a compact (4P, kc, kc) scatter plus a top-wide_row_budget full-width supplement. Output is bit-identical for every geometry within budget. On overflow (more than wide_row_budget rows wider than compact_width) the returned matrix is NaN, following the mesh/interpolator/sibson.py cap convention, rather than silently truncated. K <= compact_width callers (Delaunay mesh, K = 4) are unaffected. The NumPy path (pixel_splitted_regularization_matrix_np_from) is unchanged.

Migration

• None required.

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

https://claude.ai/code/session_011xsh8KqgiHWTfPGgWEYMQw

[Jammy2211]

Workspace PRs (merge after this one, library-first gate):

• test(jax_assertions): ConstantSplit compaction parity + per-geometry wide-row count (PyAutoArray#537) autolens_workspace_test#309 — compaction parity assertion + per-geometry wide-row count
• results: A100 A/B of the compacted ConstantSplit assembly (PyAutoArray#537, 12x) autolens_profiling#231 — assembly bench, A100 A/B results, results note