[Jammy2211]

Summary

Companion measurement-of-record for phase 2 of the gaussian-deflections-precompute epic (PyAutoGalaxy#604): the JAX trace-time fold of fixed-geometry deflection fields. Library PRs: PyAutoArray (constant grid shift-and-rotate at compile time) then PyAutoGalaxy (memo JAX branch) — both must merge first.

Headline on the SLaM-shaped hst fit under jit(vmap) (batch 3): with the memo on, the compiling call makes 180 numpy-backend and 0 JAX-backend Faddeeva calls (off: 0 / 240), steady state 0 / 0 both; jax_folds 90 (30 Gaussians × lp / pixelization / blurring grids); jaxpr 53,369 → 13,289 equations (−75%); log-likelihood agrees to 8.6e-15; vmap_first_call 10.8 → 5.4 s (2.0×), 9.96 → 6.08 s on a second run (1.64×); vmap_steady_x10 unchanged at 2.5–2.6 s — the compiled inversion dominates this fit, so the fold buys compile time and graph size here, not steady-state throughput. Controls: the kill switch reproduces memo-off exactly; a free grid_offset makes the grid a tracer, the memo falls through (jax_folds 0), value unchanged.

Finding recorded: the plan's premise that the fit grid is concrete inside jit was false — JAX stages every jnp op, so subtracted_and_rotated_from returned a tracer at all 132 memo-hook calls; the fix is PyAutoArray's ensure_compile_time_eval on a constant shift-and-rotate. Also recorded: a DatasetModel with a free grid_offset cannot round-trip autofit.jax.register_model (Prior.tree_unflatten TypeError), so the control uses the flat-vector path — pre-existing, candidate follow-up.

Scripts Changed

• scripts/imaging/likelihood_runtime/mge_mass_jax.py — new cell: SLaM-shaped model under AnalysisImaging(use_jax=True), memo-off vs memo-on legs in one process, _wofz witness split by backend, memo_stats(), jaxpr equation counts, likelihood agreement, vmap_first_call / vmap_steady_x10 (3 repeats), kill-switch and free-grid_offset controls, pinned EXPECTED_LOG_LIKELIHOOD
• results/runtime/imaging/mge_mass_jax/ — artifact
• results/notes/numpy_deflections_cpu.md — section "Fixed-geometry deflection memo — phase 2: JAX trace-time constant (2026-09-03)": design, the concreteness finding and fix, witness counts, jaxpr delta, timing table, controls, workspace_test pin table (15/15, bit-identical), findings

Upstream PR

• perf: evaluate a constant grid shift-and-rotate at compile time under JAX (memo phase 2, PyAutoGalaxy#604) PyAutoArray#520 — must merge first
• perf: fold fixed-geometry deflection fields out of the JAX trace (memo phase 2, #604) PyAutoGalaxy#605 — must merge second

Test Plan

• python scripts/imaging/likelihood_runtime/mge_mass_jax.py runs both legs and controls; pinned_drift: []
• scripts/lens/deflections/{total,dark,stellar,basis}.py hst numpy pins PASSED at rtol 1e-6
• scripts/misc/tooling/build_readme.py --check clean
• ruff check / ruff format --check clean

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

https://claude.ai/code/session_01XhnA4pFN2NycuKc8Ni6s2R