jax_ile: --angle-marg-scheme multipeak (#313) is not traceable under the sampler jit; Table 3 arms cannot run it

[oshaughnessy-junior]

Found 2026-09-09 on rift_O4d 336f861 by the Table 3 (paper 1 tab:sampler_arms) campaign.

Symptom. Every driver run with --angle-marg-scheme multipeak dies in ~30 s at the flow sampler's jitted warmup: anglemarg.py:2517 in fused_log_likelihood_distphipsimarg_multipeak converts a traced array to numpy (jax.errors.TracerArrayConversionError, complex128[3,3,2,646]), reached via wrapper._fused -> samplers._warmup_compile. Traceback: /scratch/richard.oshaughnessy/table3_arms_20260909/out/flow_snr40_multipeak_seed9101/stderr.log (CIT).

Why the tests missed it. #313's driver-seam test checks option parsing and the library test pins one row through an eager call; nothing traces the fused kernel under jax.jit. The controller was validated with eager per-row calls (64 rows in 6.5 min sequential, ~5 s per row on one Blackwell), which is also its cost model: a sampler evaluating thousands of extrinsic samples per stage cannot use it as-is even if it traced.

What is needed. (1) A test that traces the fused kernel under jax.jit on carrier tables, so a scheme cannot be registered without it. (2) Either a traceable kernel (no host numpy, Python control flow and the Newton loops as lax primitives, batched over samples) with a measured per-sample cost, or a jax.pure_callback wrapper whose per-sample cost is measured and stated, so the scheme is honestly either an in-sampler likelihood or a per-row controller. (3) Until then, the paper's Table 3 differentiable-arm rows stay red (RO, 2026-09-09: "skip flow for now; add issue; leave red in table for today").

Related: the reserve rule in #313 duplicates #304's peaklocal_time_reserve; converge on one (the #304 kernel carries the focus certificate).

[oshaughnessy-junior]

Answers to the two questions, measured rather than estimated.

1. Traceability

Not without rewriting the planner. multipeak_planner.py carries 127 lines of
host-side Python control flow and numpy/scipy calls that a trace cannot cross:
scipy.special.logsumexp, np.linalg.solve on per-mode Cholesky factors, Python
for loops over retained modes, and if branches on data-dependent values
(result0.ok, capacity_truncated, the mode-merge decisions). The clipped Newton
iteration is a Python loop with a data-dependent stop. Converting these to
lax.scan/lax.cond with static shapes is a rewrite of the planner, not a patch,
and the discrete selection it performs is the part the design deliberately keeps
outside the differentiated graph.

The sampler can be made to run without that rewrite, by not tracing the
scheme at all: leave wrapper._batched un-jitted for multipeak only. That is
what PR #317 does, and its new test crosses the _batched seam that my original
tests missed. It is the pure_callback-equivalent outcome, eager per batch,
without the callback plumbing.

2. Batched cost per extrinsic sample

There is no vmapped cost to quote: the planner is a Python loop over rows, so
per-sample cost does not amortize. Measured on one RTX PRO 4000 Blackwell, 64
rows per rung:

rho	s per sample	1000 samples	10000 samples
40.77	5.17	86 min	14.4 h
163.08	4.74	79 min	13.2 h
652.31	5.17	86 min	14.3 h

For comparison the AV arm completes a whole rung in 28 s (rho 40), 81 s (163) and
482 s (640).

So the local branch is not viable as a per-sample sampler likelihood at this
cost, independently of the tracing question. A sampler stage of a thousand
samples costs more than an hour where AV completes the rung in minutes. Fixing
traceability alone would not make Table 3 runnable; the cost would have to come
down by three or four orders of magnitude, which for a host-side planner means
the rewrite above plus batching across samples.

The controller's flat-in-rho cost (4.7-5.2 s from rho 40.77 to 652.31, ~218 MiB)
is a real result about the algorithm, recorded in
analyses/va_sequence_20260902/RESULTS_20260909_multipeak_ladder.md. It is a
statement about per-evaluation scaling, not about sampler throughput.

Related: this reserve rule and #304's peaklocal_time_reserve are two peak-local
time rules and should converge on one, probably #304's, which carries a focus
certificate mine lacks.