feat: JAX Delaunay walk — early-exit while_loop, chunk only the seed argmin

[Jammy2211]

Overview

On the A100 the HST Delaunay imaging likelihood (1500 Hilbert vertices, 15,361 data pixels, MGE-60 lens, ConstantSplit) costs 97 ms per evaluation, of which the --split-setup decomposition attributes 26.6 ms to "Triangulation + interpolation". Almost all of that is latency in the JAX point locator pix_indexes_delaunay_walk_from (autoarray/inversion/mesh/interpolator/delaunay.py): a fixed-trip fori_loop of 128 steps that cannot exit early (typical resolution < 10 steps), run over 1024-query chunks through a sequential lax.map (~22 chunks x 128 = ~2,800 dependent steps per likelihood, since jax_delaunay locates the data grid and the 6000 split-cross points separately). The chunking exists only to bound the (chunk, N) nearest-vertex argmin intermediate under vmap; it happens to serialise the latency-bound walk as well.

This issue is Phase 1 of the prompt: early-exit while_loop, chunk only the seed argmin, one walk call per likelihood, then re-measure on the A100. Phase 2 (static image-plane seed + one-shot fan test, Mapper/AdaptImages plumbing) is filed separately only if Phase 1 leaves the row above a few ms.

Witness: the "Triangulation + interpolation" row of autolens_profiling/results/breakdown/imaging/delaunay_hpc_a100_fp64.json (--split-setup) drops from 26.6 ms to under 8 ms unbatched, EXPECTED_LOG_EVIDENCE_HST unchanged, walk parity tests pass, FD certification passes.

Constraint (user): no slowdown or behaviour change for the numba CPU likelihood or the JAX CPU likelihood. The numba/NumPy path goes through scipy_delaunay (find_simplex + KDTree) and never enters the walk, so it is untouched by construction; the JAX CPU path is gated by a local before/after breakdown run.

Plan

1. Replace the fixed-trip fori_loop in the JAX walk with a while_loop that exits as soon as every query is located or outside the hull; keep the 128-step cap as the safety bound.
2. Split the chunk loop's two jobs: keep only the nearest-vertex argmin chunked (memory guard), then run the walk once over all queries with no lax.map.
3. Locate data-grid and split-cross points in a single walk call in jax_delaunay (concatenate, locate, slice) so one while_loop runs per likelihood instead of two.
4. Freeze the walk's float inputs with stop_gradient (point location is integer-valued; while_loop has no reverse-mode rule); barycentric weights are still recomputed from traced points, so jax.grad through the likelihood is unchanged.
5. Keep the NumPy path (unit tests) behaviourally identical; extend the NumPy unit tests to the refactored seed/walk split.
6. Prove JAX parity against scipy.spatial.Delaunay.find_simplex in a workspace_test jax_assertions script (existing meshes + a 1500-vertex Hilbert mesh traced through a few mass models), plus jit and vmap round-trips; re-run the FD certification.
7. Ship the library PR; re-measure on the A100 from a RAL worktree, record the four-way split, ship the profiling/workspace_test PR. Decide Phase 2 from the numbers.

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary, library)
• autolens_workspace_test (JAX parity script — unit tests stay NumPy-only)
• autolens_profiling (A100 re-measurement results + notes)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean (orphan worktree delaunay-area-magnification-audit on a feature branch — unregistered, not a conflict)
./autolens_workspace_test	main	clean
./autolens_profiling	main	clean (claimed by retire-gpu1-mig-exclusion, disjoint file set — parallel claim)

Suggested branch: feature/delaunay-walk-early-exit

Implementation Steps

A. PyAutoArray — autoarray/inversion/mesh/interpolator/delaunay.py

Refactor pix_indexes_delaunay_walk_from into three pieces, public signature unchanged (callers: sibson.py x2 with return_simplex_indexes=True, test_delaunay_walk.py, autolens_workspace_test/scripts/misc/jax_assertions/delaunay_nn.py):

1. _nearest_vertex_seed_from(query_points, points, xp) → (Q,) int32. NumPy: one brute-force argmin as today. JAX: pad Q to a multiple of DELAUNAY_LOCATE_CHUNK, jax.lax.map over (-1, chunk, 2) computing only the per-chunk argmin, reshape, slice to Q. Pad rows at 1e9 as now. Sole survivor of the chunking; the (chunk, N) intermediate stays bounded under vmap.
2. _walk_from_seed(query_points, seed, points, simplices_padded, simplex_neighbors, vertex_simplex, xp) → (cur, done, outside) over all Q at once. walk_step / weights_of unchanged in content. NumPy: existing for … if (done | outside).all(): break. JAX: jax.lax.while_loop(cond, body, (step, cur, done, outside)) with cond = (step < DELAUNAY_WALK_STEPS) & jnp.any(~done & ~outside). Under vmap the loop runs to the slowest lane.
3. Wrapper builds mappings / simplex_indexes from (seed, cur, done) exactly as the current locate_chunk tail does, honouring return_simplex_indexes.

Gradient contract: lax.while_loop is not reverse-mode differentiable, so the JAX branch wraps query_points and points in jax.lax.stop_gradient before the seed argmin and the walk (the connectivity tables already are, see _jax_delaunay_tables). Location is integer-valued with zero a.e. derivative; every differentiable quantity (barycentric weights via pixel_weights_delaunay_from, dual areas, split points, Sibson weights) is computed downstream from the traced arrays. Document in the docstring; update the DELAUNAY_LOCATE_CHUNK / DELAUNAY_WALK_STEPS comments to the new roles.

jax_delaunay: locate jnp.concatenate([query_points, split_points]) in one call and slice [:Q] / [Q:]; split points still seed from their own nearest vertex. scipy_delaunay (NumPy path) untouched; sibson.py needs no change.

B. PyAutoArray unit tests (NumPy only, never import jax) — test_autoarray/inversion/pixelization/interpolator/test_delaunay_walk.py: existing _assert_matches_find_simplex tests unchanged; add _nearest_vertex_seed_from vs cKDTree on the blob-ring mesh (ties excluded); _walk_from_seed from a deliberately far seed (vertex 0 for every query) still reaches the find_simplex answer within the cap; wrapper returns (Q, 3) int32 and (Q,) int32 with return_simplex_indexes=True on an odd Q.

C. autolens_workspace_test — scripts/misc/jax_assertions/delaunay_walk.py (sibling of delaunay_nn.py, same __Env__ block and ENV: jax full_datasets declaration, bare """ opener; CI runs every script): parity vs find_simplex + cKDTree fallback on uniform / blob-ring meshes (reuse adaptive_mesh) and on a 1500-vertex Hilbert mesh traced through a handful of the mass models from delaunay_nn_caps.py (lives in this repo, not autolens_profiling), data grid and 4N split points, comparing mapping matrices not triangle ids on shared-edge ties; jax.jit equals eager; jitted vmap over a small batch equals per-member; odd Q; jax.grad of a scalar built from the barycentric weights runs and matches FD; print warm per-call time.

D. FD certification + CPU no-regression gate (before the library PR)

• Numba/NumPy path: grep confirms no NumPy-path call site is added; run autolens_profiling/scripts/imaging/likelihood_breakdown/delaunay_numba.py locally before/after as a control (same numbers to noise, same pin).
• JAX CPU path: run autolens_profiling/scripts/imaging/likelihood_breakdown/delaunay.py --split-setup locally on main and on the feature branch, warm, three repeats each. Gate: "Triangulation + interpolation" row and total per-eval time not slower; EXPECTED_LOG_EVIDENCE_HST unchanged. Both tables go in the PR.
• autolens_workspace_test/scripts/imaging/jax_grad/delaunay.py must pass.
• Scratchpad micro-benchmark of the locator alone (N=1500, Q=15,361 + 6000 split points) on CPU, before vs after.
• Read the cloud#34099198772 failure for autolens_test scripts/imaging/delaunay.py and state whether it touches the walk.

E. Ship + A100 verification

1. ship_library: tests, commit, push, PR on PyAutoArray (pending-release).
2. From a RAL worktree with HPCPullPyAuto on the feature branch, run scripts/imaging/likelihood_breakdown/delaunay.py --config-name hpc_a100_fp64 --split-setup --vmap-batch 16 and the delaunay_nn.py sibling. Pins must pass unchanged — a shift is a bug, not a re-pin (PyAutoLens#721 knife-edge lesson: bisect before re-pinning anything).
3. Results JSON/PNG under autolens_profiling/results/breakdown/imaging/ and results/notes/delaunay_walk_early_exit.md with the four-way split unbatched and per call under vmap, the H row and the single-JIT cell against the 2026-09 baselines.
4. ship_workspace for autolens_workspace_test + autolens_profiling behind the library-first merge gate. Decide whether to file the Phase 2 prompt.

Key Files

• PyAutoArray/autoarray/inversion/mesh/interpolator/delaunay.py — pix_indexes_delaunay_walk_from, jax_delaunay, _jax_delaunay_tables, the two module constants
• PyAutoArray/autoarray/inversion/mesh/interpolator/sibson.py — DelaunayNN caller (unchanged; inherits the gain)
• PyAutoArray/test_autoarray/inversion/pixelization/interpolator/test_delaunay_walk.py — NumPy parity tests
• autolens_workspace_test/scripts/misc/jax_assertions/delaunay_walk.py — new JAX parity/jit/vmap/grad script
• autolens_workspace_test/scripts/misc/jax_assertions/delaunay_nn_caps.py — lensing geometry to reuse
• autolens_workspace_test/scripts/imaging/jax_grad/delaunay.py — FD certification
• autolens_profiling/scripts/imaging/likelihood_breakdown/delaunay.py, delaunay_nn.py, delaunay_numba.py — breakdown / control scripts
• autolens_profiling/results/notes/preopt_breakdown_baseline.md — the 26.6 ms baseline

Original Prompt

Click to expand starting prompt

JAX Delaunay point location: early-exit walk, unchunked loop, static image-plane seed

Type: feature
Target: autoarray
Repos:

• PyAutoArray
• autolens_profiling
  Themes:
• jax-gpu
• delaunay
• profiling
• performance
  Difficulty: medium
  Autonomy: supervised
  Priority: high
  Status: draft
  Consequence: judge
  Witness: the A100 Delaunay breakdown's "Triangulation + interpolation" row (26.6 ms, results/breakdown/imaging/delaunay_hpc_a100_fp64.json, --split-setup) drops to under 8 ms unbatched with EXPECTED_LOG_EVIDENCE_HST unchanged, the walk parity tests pass, and the FD certification still passes
  Review-minutes: 30
  Unattended: ready
  Filed: 2026-09-05

Original request (verbatim):

give me a prompt to work on this Target 1: point location is a fixed-trip loop over chunks

The measurement

On the A100 the HST Delaunay imaging likelihood (1500 Hilbert vertices, 15,361 data
pixels, MGE-60 lens, ConstantSplit) costs 97 ms per evaluation, of which the four-way
--split-setup decomposition attributes 26.6 ms to "Triangulation + interpolation"
(autolens_profiling/results/notes/preopt_breakdown_baseline.md). The qhull host callback
is a few ms of that at most (1500 points, pure_callback, vmap_method="sequential"). The
rest is the JAX-side point location in
PyAutoArray/autoarray/inversion/mesh/interpolator/delaunay.py,
pix_indexes_delaunay_walk_from, and it is latency, not FLOPs:

• The walk is jax.lax.fori_loop(0, DELAUNAY_WALK_STEPS, ...) with
  DELAUNAY_WALK_STEPS = 128 (module constant at line 75, loop at line 252). A fixed-trip
  fori_loop cannot exit early. Seeded from the nearest vertex the walk resolves in a
  handful of steps, so well over 95% of the 128 iterations are no-ops that still launch a
  gather, a cross-product batch, an argmin and three wheres each.
• Queries are processed in chunks of DELAUNAY_LOCATE_CHUNK = 1024 (line 69) through
  jax.lax.map (line 279), which is sequential: 16 chunks for the data grid and 6 more for
  the 6000 ConstantSplit split-cross points (jax_delaunay calls the walk twice). That is
  roughly 22 x 128 = 2,800 dependent walk steps per likelihood.
• The chunking exists only to bound the (chunk, N) nearest-vertex distance intermediate
  under vmap (184 MB per replica at full Q; ~12 GB at batch 64). It is a memory guard
  that happens to serialise the latency-bound part.

Nautilus at n_batch=256 gave 51 ms per eval against 62 ms at n_batch=16, so batching
barely amortises this today. DelaunayNN (interpolator/sibson.py) seeds its cavity walk
from pix_indexes_delaunay_walk_from(..., return_simplex_indexes=True), so it inherits
every gain here. interpolator/knn.py is a separate approach and out of scope.

Phase 1: early exit, and chunk only the argmin

1. Replace the fori_loop with a jax.lax.while_loop whose predicate is
   (step < DELAUNAY_WALK_STEPS) & jnp.any(~done & ~outside). The 128 cap stays as the
   safety bound; the typical exit is under ten steps. Under vmap the loop runs to the
   slowest lane, which is fine.
2. Split the two jobs the chunk loop currently does. Keep the nearest-vertex argmin
   chunked (it is the memory hazard, and a lax.map over 16 single-argmin chunks is cheap),
   then run the walk once over all Q queries with no lax.map. The walk's working set is
   O(Q), not O(Q x N).
3. Keep the NumPy path (xp is np, used by the unit tests) behaviourally identical; it
   already early-exits.
4. Prove parity: the JAX walk must return the same (Q, 3) mappings and simplex indexes as
   scipy.spatial.Delaunay.find_simplex on the existing test meshes and on a 1500-vertex
   Hilbert mesh traced through a few mass models (reuse the geometry from
   autolens_profiling/scripts/misc/jax_assertions/delaunay_nn_caps.py). Points exactly on
   a shared edge may legitimately resolve to either adjacent triangle; the barycentric
   weights agree because the opposite vertex gets weight zero, so compare mapping matrices,
   not triangle ids, in that case.

Phase 2: static image-plane seed and one-shot fan test

The image-plane mesh is fixed per fit by the adapt image, and ray tracing is continuous
away from critical curves, so each data pixel's nearest image-plane mesh vertex is a
near-perfect source-plane seed. Precomputing that index once per fit (in the Mapper /
AdaptImages layer, where image_plane_mesh_grid is known) removes the brute-force
(Q, N) argmin and its memory hazard entirely, which also retires the chunking. Then have
the qhull callback also return a padded vertex-to-incident-triangle table (cap around 12,
audit the actual max the way delaunay_nn_cap_audit.md did) so most queries resolve with
one vectorised barycentric test over the seed's fan; the while_loop walk from Phase 1
handles only the residual near critical curves and outside the hull. Split-cross points
have no image-plane parent pixel; seed them at their parent vertex (they are offsets from
it) and let the fan test cover them.

Phase 2 changes the interpolator's inputs (a seed array), so it touches
InterpolatorDelaunay, InterpolatorDelaunayNN, the Mapper constructor path and
FitImaging plumbing. Land Phase 1 first and re-measure; Phase 2 is only worth its
plumbing if Phase 1 leaves the row above a few ms.

Gradient contract (do not relax)

Point location is integer-valued; its derivative is zero almost everywhere and it is
already wrapped by stop_gradient semantics through the frozen connectivity tables
(_jax_delaunay_tables docstring). The barycentric weights are recomputed from the traced
points after location exactly as now, so jax.grad through the likelihood is unchanged.
Re-run the FD certification autolens_workspace_test/scripts/imaging/jax_grad/delaunay.py
after each phase.

Verification on the A100

Use the delaunay-nn-breakdown tooling (autolens_profiling#219): rerun
scripts/imaging/likelihood_breakdown/delaunay.py --config-name hpc_a100_fp64 --split-setup --vmap-batch 16 and the delaunay_nn.py sibling from a RAL worktree with
HPCPullPyAuto pointed at the feature branch. Report the four-way split unbatched and per
call under vmap, the H row, and the single-JIT runtime cell, against the 2026-09 baselines.
Pins (EXPECTED_LOG_EVIDENCE_HST in both scripts) must pass unchanged; a shift means a
mapping changed and is a bug, not a re-pin. Note the symmetric knife-edge lesson from
PyAutoLens#721: if a positions-threshold or point-solver pin elsewhere moves, bisect
before re-pinning.

[Jammy2211]

Library PR Created

PR: #531 (pending-release, commit 6152bdf0, base bcd15cd9)

Next: /start_workspace — the workspace follow-up (new parity script + A100 re-measure) is a separate PR and is not required for #531 to merge.

What shipped

pix_indexes_delaunay_walk_from is split into _nearest_vertex_seed_from (the chunked argmin — now the only chunked stage, still the memory guard for the (chunk, N) intermediate under vmap) and _walk_from_seed (a lax.while_loop over all Q queries at once, cond (step < 128) & any(~done & ~outside)), behind a wrapper whose public signature and return_simplex_indexes contract are unchanged. DELAUNAY_WALK_STEPS is now a safety cap, not a trip count. The JAX branch stop_gradients query_points/points (while_loop has no reverse-mode rule; the locator returns only integer indices, whose true derivative is zero away from measure-zero re-wiring events; every differentiable downstream quantity is recomputed from the traced arrays). jax_delaunay now locates the data grid and the split-cross points in one concatenated call. sibson.py (DelaunayNN) inherits the gain; the NumPy path (scipy_delaunay, numba CPU likelihood) is untouched.

Verification

• pytest test_autoarray/ — 1452 passed; 3 new NumPy-only tests in test_delaunay_walk.py (seed vs cKDTree.query; walk converges to find_simplex from a far seed; wrapper shapes/dtypes on an odd Q).
• JAX parity vs scipy find_simplex + cKDTree fallback: exact-row fraction 1.00000 on uniform N=400/Q=3000, blob-ring N=400/Q=3000, blob-ring N=1500/Q=15974; jit == eager; odd Q; jit(vmap) batch == per-member; jax.grad runs, central-FD rel err 1.8e-11 … 1.6e-9.
• FD certification autolens_workspace_test/scripts/imaging/jax_grad/delaunay.py — PASSED (exit 0; mass/shear rel err 7.4e-6 … 1.1e-3 under rtol 1e-2).
• CPU no-regression gate: JAX CPU likelihood_breakdown/delaunay.py --split-setup, median of 5 warm reps — params→H prefix 179.31 → 106.45 ms (1.68x), every after-rep below every before-rep; TOTAL 2932.39 → 2797.60 ms; EXPECTED_LOG_EVIDENCE_HST = 29110.92085793 PASSED 5/5 both sides, not re-pinned. Numba control delaunay_numba.py: TOTAL 0.325 → 0.297 s, pin PASSED (never enters the walk).
• Micro-benchmark (CPU fp64 jitted, median of 10): locator Q=21361 124.15 → 76.70 ms (1.62x); jax_delaunay 131.96 → 68.94 ms (1.91x).

⚠️ Witness caveat — the A100 verification (plan section E) must be judged on the params→H prefix

Before this change the breakdown script's "Triangulation + interpolation" prefix returned only the data-grid mappings, so XLA dead-code-eliminated the separate split-point walk and that cost was charged to the H row. With one concatenated walk the split walk is now live inside that prefix, so the row is flat (118.92 → 120.11 ms) and the H row goes negative (55.58 → −13.48 ms). Neither row is interpretable on its own any more.

The honest apples-to-apples witness is the params→H prefix (equivalently Tri+interp + H). The PyAutoMind task's Witness: line has been amended to that measure; the profiling script's H-row attribution is a follow-up for the workspace PR.

Readiness

Heart at ship time: YELLOW, no RED reasons — "workspace validation not passing (5 failed, 2 timeout, cloud#34099198772: …)" plus stale "release validation incomplete: no rehearsal for current source". Organism-scope: the failing autolens_test scripts/imaging/delaunay.py legs in that run do not exercise the walk and were already fixed on autolens_workspace_test main today (078e445, 4103234).

[Jammy2211]

Workspace PR opened — 2026-09-07

• Workspace PR: delaunay walk: JAX assertion gate for the early-exit locator (PyAutoArray #530) autolens_workspace_test#306 (pending-release)
• Library PR: feat: JAX Delaunay walk — early-exit while_loop, chunk only the seed argmin #531

autolens_workspace_test gains scripts/misc/jax_assertions/delaunay_walk.py, the correctness gate for the early-exit lax.while_loop locator, registered in smoke_tests.txt beside the delaunay_nn.py siblings so it runs on the PR and weekly/integrate channels. It asserts exact parity against scipy.spatial.Delaunay.find_simplex + cKDTree on uniform, blob-ring and production-like 1,500-vertex lensing meshes (data grid and 4N split points, shared-edge ties only and a tie must still contain the query); that jax.jit equals eager on a query count past the 1,024-wide seed-chunk boundary and on an odd short set, so the chunked argmin's padding never leaks; that a jitted vmap equals the per-member results, since the while_loop runs to the slowest lane in a batch; and that jax.grad through the barycentric weights runs at all (the locator stop_gradient-wraps its float inputs because lax.while_loop has no reverse-mode rule) and matches central finite differences. Warm locator time is printed, not gated. Validated against the worktree autoarray carrying the #531 change: 17/17 PASS in 12.8 s direct and PASS in 12.4 s through the shared run_python smoke runner, black --check clean. The script imports the rewritten locator, so the workspace PR is held behind the library-first gate — merge #531 and ship its release first.

[Jammy2211]

A100 verification — 2026-09-07

Witness verdict: PASSED, by 2.1×. Same-node, same-session A/B on
euclid-ral-gpu-2, jobs 342315–342320, all six in one 14-minute window. Control =
private PyAutoArray bcd15cd9 checkout (the branch's merge base), feature = 6152bdf0
(feature/delaunay-walk-early-exit); the shared /mnt/ral/jnightin/PyAuto install was
deliberately not touched — the subhalo-validation and Euclid DR1 CPU runs were live on it
— so each leg prepends its own PyAutoArray checkout to PYTHONPATH and prints
autoarray.__file__ + git rev-parse HEAD into its SLURM log.

Delaunay (HST / Hilbert-1500 / MGE-60 / ConstantSplit, fp64), ms per call

Row	control	feature	Δ
Triangulation + interpolation	30.563	5.970	−24.593
Regularization matrix (H)	13.555	0.225	−13.330
params→H prefix (the witness)	45.149	7.228	−37.921 (6.25×)
Total step-by-step (unbatched)	99.463	62.045	−37.418
Full pipeline, single JIT	93.463	60.355	−33.108 (1.55×)

The witness asked for a ≥ 18 ms unbatched drop in the params→H prefix; measured
−37.9 ms.

Under vmap 16 (per call): params→H prefix 7.555 → 5.139 ms, full pipeline
41.945 → 39.520 ms. The batched saving is small because at batch 16 the walk was
already amortising and the dense linear algebra dominates.

DelaunayNN sibling — shares the walk but not the concatenation
(sibson.py:522 / :648 still call pix_indexes_delaunay_walk_from twice): params→H
prefix 178.971 → 144.789 ms (1.24×), 28.757 → 24.424 ms per call at vmap 16.

Pins

EXPECTED_LOG_EVIDENCE_HST unchanged on every leg — 29110.920858 for Delaunay
(342315/342316, pinned_drift: [] on both runtime rows) and 29144.581944 for
DelaunayNN (342317/342318). The mapping is bit-identical to the eager NumPy reference on
both sides of the change.

H-row attribution caveat

The H row moved meaning with this change; the params→H prefix did not. Step 11 is
attributed as t(params → H) − t(params → interpolator outputs); before this PR a prefix
stopping at step 6 never asked for the ConstantSplit split points, so XLA
dead-code-eliminated that second walk out of the interpolator prefix and its whole cost
landed in the H row's subtraction. With the concatenated locate the split walk runs
inside prefix 6, so the subtraction leaves only the regularization assembly — hence
H: 13.555 → 0.225 ms, and a small negative would have been equally legitimate. Compare
regularization_matrix_prefix_s (params→H) across library versions, not either row
alone. Recorded in the breakdown script's docstring in the profiling PR. The
NumPy/scipy path is untouched and still charges the split walk to H.

Phase 2 recommendation — do not start it

Do not start Phase 2 (static image-plane seed + fan test) on the Delaunay path. After
this PR the Delaunay params→H prefix is 7.23 ms unbatched and 5.14 ms per call at
vmap 16 — ~13% of the 39.5 ms batched per-evaluation cost, so Phase 2's remaining
headroom on this cell is a few ms per call for a change that reaches into
Mapper/AdaptImages plumbing. DelaunayNN still spends 144.8 ms unbatched / 24.4 ms per
call in the same prefix, and the cheapest cut there is the one Delaunay already took:
concatenate sibson.py's two walk calls into one. Filed instead as
PyAutoMind/draft/feature/autoarray/sibson_single_concatenated_walk.md.

Links

• Library PR: feat: JAX Delaunay walk — early-exit while_loop, chunk only the seed argmin #531
• Workspace-test PR: delaunay walk: JAX assertion gate for the early-exit locator (PyAutoArray #530) autolens_workspace_test#306
• Profiling results PR: A100 re-measure: JAX Delaunay walk early exit (PyAutoArray#531) autolens_profiling#224
• Full note (source of truth for every number above):
  results/notes/delaunay_walk_early_exit.md in autolens_profiling#224

[Jammy2211]

Shipped

All three PRs merged 2026-09-07; Phase 1 of the prompt is complete.

• Library: feat: JAX Delaunay walk — early-exit while_loop, chunk only the seed argmin #531 (pending-release — PyAutoArray has not yet published a release carrying this change)
• Workspace parity gate: delaunay walk: JAX assertion gate for the early-exit locator (PyAutoArray #530) autolens_workspace_test#306
• Profiling results: A100 re-measure: JAX Delaunay walk early exit (PyAutoArray#531) autolens_profiling#224

What shipped

pix_indexes_delaunay_walk_from is split into _nearest_vertex_seed_from (the chunked nearest-vertex argmin, now the only chunked stage and still the memory guard for the (chunk, N) intermediate under vmap) and _walk_from_seed (a lax.while_loop over all Q queries at once, cond (step < DELAUNAY_WALK_STEPS) & any(~done & ~outside)), behind a wrapper whose public signature and return_simplex_indexes contract are unchanged. DELAUNAY_WALK_STEPS is now a safety cap rather than a trip count. jax_delaunay locates the data grid and the split-cross points in one concatenated call, so one walk runs per likelihood instead of two. The JAX branch stop_gradients its float inputs (lax.while_loop has no reverse-mode rule; the locator returns integer indices only, and every differentiable downstream quantity is recomputed from the traced arrays). The NumPy/scipy_delaunay path — and therefore the numba CPU likelihood — is untouched by construction.

Witness verdict: PASSED, by 2.1x

Same-node A/B on euclid-ral-gpu-2 (jobs 342315–342320, one 14-minute window), control bcd15cd9 vs feature 6152bdf0, HST / Hilbert-1500 / MGE-60 / ConstantSplit, fp64:

Row	control	feature
params→H prefix (the witness)	45.149 ms	7.228 ms (6.25x)
Full pipeline, single JIT	93.463 ms	60.355 ms (1.55x)

The witness asked for a ≥ 18 ms unbatched drop in the params→H prefix; measured −37.9 ms. At vmap 16 the prefix is 7.555 → 5.139 ms. EXPECTED_LOG_EVIDENCE_HST unchanged on every leg (29110.920858 Delaunay, 29144.581944 DelaunayNN), no pin re-pinned.

CPU no-regression gate (the user constraint): JAX CPU params→H prefix 179.31 → 106.45 ms, every after-rep below every before-rep; numba control delaunay_numba.py TOTAL 0.325 → 0.297 s with the pin passing — it never enters the walk.

Verification

• pytest test_autoarray/ — 1452 passed, including 3 new NumPy-only tests in test_delaunay_walk.py.
• scripts/misc/jax_assertions/delaunay_walk.py (new, registered in smoke_tests.txt) — 17/17 PASS: exact parity vs scipy find_simplex + cKDTree on uniform, blob-ring and 1,500-vertex lensing meshes, jit == eager past the 1,024 seed-chunk boundary and on odd Q, jit(vmap) == per-member, jax.grad runs and matches central FD.
• FD certification scripts/imaging/jax_grad/delaunay.py — PASSED.

H-row attribution caveat

With one concatenated locate the split walk now runs inside the interpolator prefix, so the "Triangulation + interpolation" and "Regularization matrix (H)" rows are no longer individually comparable across this change (H goes 13.555 → 0.225 ms purely by re-attribution). Compare regularization_matrix_prefix_s (params→H) across library versions, not either row alone. Recorded in the breakdown script's docstring.

Phase 2 — deliberately not filed

Phase 2 of the prompt (static image-plane seed + one-shot fan test, Mapper/AdaptImages plumbing) is not being filed. After this PR the Delaunay params→H prefix is 7.23 ms unbatched and 5.14 ms per call at vmap 16 (~13% of the 39.5 ms batched per-evaluation cost), so its remaining headroom is a few ms for a change that reaches into Mapper/AdaptImages plumbing. The cheaper cut is the one Delaunay already took, applied to DelaunayNN (still 144.8 ms unbatched / 24.4 ms at vmap 16 in the same prefix, because sibson.py still calls the locator twice) — filed instead as PyAutoMind/draft/feature/autoarray/sibson_single_concatenated_walk.md.

Full note, the source of truth for every number: results/notes/delaunay_walk_early_exit.md in autolens_profiling.