feat: MultiStartGradient value/gradient NaN step diagnostics

[Jammy2211]

Overview

MultiStartGradient detects a dead lane only via the value (alive = np.isfinite(np.asarray(foms))). The gradient is never checked inside the fit loop, so a lane whose likelihood is finite but whose gradient is non-finite is counted alive, has its update zeroed by optax.apply_if_finite, and silently freezes in place. That failure mode — a differentiability failure rather than an evaluation failure — is currently invisible.

This adds per-step, disjoint accounting for both: value-NaN lane-steps (likelihood undefined) and gradient-NaN lane-steps (likelihood defined but not differentiable), persisted into search_internal and samples_info and surfaced in search.summary as raw counts plus rates normalised by n_starts × total_steps.

Scope is measurement only — resurrect behaviour is unchanged, so existing benchmark numbers stay comparable. A standalone demonstration in autolens_profiling shows the accounting does not cost runtime.

Plan

• Add two per-step lane-death counters to the MultiStartGradient fit loop: value-NaN and gradient-NaN lane-steps.
• Keep the counters disjoint — gradient-NaN counts only lanes still alive by value, so a lane non-finite in both is one value-NaN, not one of each.
• Compute gradient finiteness inside the existing jitted call (returned as a third output), so there is no extra kernel dispatch and no extra host round-trip.
• Count unconditionally, not under if self.resurrect — with resurrect=False the value-NaN counter is the only record that lanes died.
• Persist both into search_internal (resume-safe) and samples_info, alongside n_resurrections.
• Emit a guarded block in search.summary with raw counts and normalised rates, following the existing MCMC duck-typed precedent.
• Extract the per-step accounting into a pure-NumPy helper so it is unit-testable without JAX.
• Add a standalone overhead demonstration to autolens_profiling, with a duplicate-baseline control so the measurement noise floor is reported alongside the overhead.

Detailed implementation plan

Work Classification

Library (PyAutoFit) + profiling artifact (autolens_profiling) — library-led; both repos share one worktree and one branch. Two PRs; the profiling PR merges second, since its numbers describe the merged code.

Affected Repositories

• PyAutoFit (primary) — the feature
• autolens_profiling — the standalone runtime demonstration

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./autolens_profiling	main	clean (untracked results/ artifacts only)

Suggested branch: feature/multistart-nan-step-diagnostics

Worktree root: ~/Code/PyAutoLabs-wt/multistart-nan-step-diagnostics/

worktree_check_conflict multistart-nan-step-diagnostics PyAutoFit autolens_profiling returns clean (exit 0).

⚠️ Two unregistered worktrees on autolens_profiling, neither in active.md:

• feature/numerical-hazard-profiling — 0 commits ahead of main (fully merged); stale leftover, safe to ignore, worth a separate cleanup.
• feature/point-source-defaults-campaign — 8 commits ahead; genuinely unmerged work, but scoped to point-source defaults and touching nothing here.

Neither is a conflict; recorded so a later session does not rediscover them.

Implementation Steps

1. autofit/non_linear/search/mle/multi_start_gradient/search.py — the counters

• Add a pure-NumPy static helper _nan_lane_counts(foms, grad_finite) -> (n_value_nan, n_grad_nan):
  • alive = np.isfinite(np.asarray(foms))
  • n_value_nan = int((~alive).sum())
  • n_grad_nan = int((alive & ~grad_finite).sum()) — disjoint by construction
  • Returns alive too, so the existing resurrection path reuses it rather than recomputing.
• Extend batched_value_and_grad (~line 495-530) to return a third output, jnp.all(jnp.isfinite(grads), axis=1), an (n_starts,) bool:
  • unbatched path: fold into the jax.jit(jax.vmap(...)) wrapper;
  • chunked path: concatenate the third output per chunk alongside foms/grads, discarding padded rows the same way.
• In the fit loop (~line 625), consume the third output and accumulate n_value_nan_lane_steps / n_grad_nan_lane_steps. Counted unconditionally, outside the if self.resurrect guard.
• Initialise both to 0 on the fresh path (~line 574); restore via int(search_internal.get(..., 0)) on the resume path (~line 549), matching the n_resurrections precedent.
• Add both keys to the search_internal dict (~line 700).

2. search.py — samples_via_internal_from (~line 935)

Add both counters to samples_info via .get(..., 0) so pre-existing search_internal files load without a KeyError.

3. autofit/text/text_util.py — search_summary_from_samples (line 115)

Add a guarded block following the hasattr(samples, "total_accepted_samples") idiom three lines above, keyed off getattr(samples, "samples_info", {}) with .get():

Resurrections = 797
Value-NaN Lane-Steps = 797
Gradient-NaN Lane-Steps = 12
Value-NaN Lane-Step Rate = 0.0166
Gradient-NaN Lane-Step Rate = 0.00025

Rates divide by n_starts * total_steps, guarded against a zero denominator. Raw counts are not comparable across runs (797 on a 16x3000 run vs 10 on an 8x300 run differ 80x raw but ~2x by rate).

Naming constraint: neutral factual counts only. Do not label these a smoothness metric in user-facing output — the resurrection-rate to HMC-divergence-rate correlation is unvalidated (separate item, wsdev#117).

4. Tests

test_autofit/non_linear/search/mle/test_multi_start_gradient.py is deliberately JAX-free (see its own comment: _fit needs jax + optax). The helper extraction in step 1 is what keeps it that way.

• _nan_lane_counts tested directly with NumPy arrays: all-finite; value-NaN only; grad-NaN only; both non-finite on the same lane (asserts disjointness).
• samples_info plumbing tested by hand-building a search_internal dict — the pattern already used by test__samples_via_internal_from and test__samples_info__stop_reason_max_steps_and_legacy_search_internal — including a legacy dict with neither key, asserting 0.
• test_autofit/text/test_text_util.py: summary block with a stub samples object, plus a negative test that a samples object without the keys emits no block.

5. autolens_profiling — standalone overhead demonstration

scripts/misc/searches/multi_start_nan_accounting_overhead.py

The feature is unconditional, so there is no flag to A/B inside one process. Instead the script measures the two halves separately and divides:

1. Denominator — time jax.jit(jax.vmap(jax.value_and_grad(...))) on a real PyAutoLens likelihood at realistic (n_starts, n_params) shapes. That is the loop's per-step cost, same code path, same shapes.
2. Numerator — the accounting variants (host pull / eager device reduce / fused-in-jit) applied to that same real gradient output.
3. Control — a duplicate baseline run, so the noise floor prints alongside the overhead.

The control is required, not decorative: without it the script can only report "we measured no impact", which is unfalsifiable. With it, it reports "overhead is below a noise floor of X%", a claim that could have failed.

Conventions: run from repo root (ruff.toml sentinel walk-up); honour AUTOLENS_PROFILING_SMOKE=1 at module top; write a version-stamped JSON + PNG pair under results/searches/multi_start_nan_accounting/ per results/README.md; ruff check + ruff format --check clean (the PR gate); short pointer added to scripts/misc/searches/README.md. Add to the hand-maintained 5-script smoke list in .github/workflows/lint.yml.

Cross-reference scripts/misc/hazards/checks/nonfinite_gradient.py — an existing detector for non-finite gradients on likelihood surfaces. Related motivation, different artifact; not a duplicate.

Prior measurement

A scratch microbenchmark (CPU, n_starts=16, n_params=30, best-of-7, duplicate-baseline control) already established the shape of the answer and drove the design:

variant	cheap obj (180 us/step)	expensive obj (1.9 ms/step)
noise floor	~1%	~1.5%
host pull of grads	+6.4%	+0.9% (below noise)
eager jnp.all(...) outside the jit	+22.6%	+3.3%
fused into the jitted call	+5.2%	+0.05% (below noise)

The eager-outside-the-jit variant is the worst of the three — it pays an un-jitted kernel dispatch plus a second host round-trip — which is why the plan fuses into the jitted call instead. Note the CPU benchmark structurally understates the host-pull variant, since on CPU "device to host" is a memcpy in the same address space; on GPU it is a real round-trip. The fused variant sidesteps the question on either backend.

The autolens_profiling script supersedes this scratch run with real likelihood shapes, and should also be run under the GPU profile.

Key Files

• autofit/non_linear/search/mle/multi_start_gradient/search.py — fit loop, batched_value_and_grad, search_internal, samples_via_internal_from
• autofit/text/text_util.py — search_summary_from_samples (line 115)
• test_autofit/non_linear/search/mle/test_multi_start_gradient.py — JAX-free unit tests
• test_autofit/text/test_text_util.py — summary block tests
• autolens_profiling/scripts/misc/searches/multi_start_nan_accounting_overhead.py — the demonstration

Deliberately out of scope

Making resurrect trigger on non-finite gradients. That would change search behaviour and shift every existing benchmark number, so the wsdev #117/#125 pixelized results would stop being comparable without re-running. The resurrection policy is decided after the counters show how often frozen lanes actually occur.

Original Prompt

Click to expand starting prompt

# MultiStartGradient value-NaN and gradient-NaN step diagnostics

Type: feature
Target: PyAutoFit
Repos:
- PyAutoFit
Difficulty: small
Autonomy: supervised
Priority: normal
Status: formalised

MultiStartGradient value-NaN and gradient-NaN step diagnostics.

Add per-step non-finite accounting to MultiStartGradient and surface it. Scope is MEASUREMENT ONLY — no change to resurrection behaviour (see 'Deliberately out of scope' below).

PART 1 (the key diagnostic) — split and record value-NaN vs gradient-NaN steps.
In autofit/non_linear/search/mle/multi_start_gradient/search.py the fit loop detects a dead lane ONLY via 'alive = np.isfinite(np.asarray(foms))' (~line 627) — the VALUE. The gradient is never checked inside the loop; gradient finiteness is tested only when drawing initial starts (~line 876). So a lane whose value is finite but whose gradient is non-finite is not counted dead and not resurrected: optax.apply_if_finite zeroes its update and the lane silently freezes in place while still counted alive. That failure mode is currently invisible, and it is exactly a differentiability failure rather than an evaluation failure. Count both, per step, and record them separately:
  - value-NaN lane-steps: where the likelihood is UNDEFINED (today's resurrection trigger)
  - gradient-NaN lane-steps: where the likelihood is defined but NOT DIFFERENTIABLE (new; currently unmeasured)
'grads' is already in hand at that point in the loop. Persist both counters into search_internal alongside n_resurrections, and into samples_info in samples_via_internal_from. Measure, do not assume, the cost: pulling grads to host is a larger device-to-host transfer than foms (n_starts x n_params vs n_starts) and sits inside the stepped loop — the step already syncs on foms, so it is likely in the noise, but benchmark it.

PART 2 — surface the counters in search.summary.
autofit/text/text_util.py:115 search_summary_from_samples(samples) already receives the samples object, and samples_via_internal_from already puts n_resurrections, n_starts, n_steps, total_steps and resurrect into samples_info. Add a guarded block emitting the resurrection count, the two NaN counters, and NORMALIZED rates (divide by n_starts * total_steps — raw counts are not comparable across runs: 797 on a 16x3000 run vs 10 on an 8x300 run differ 80x raw and ~2x by rate). Follow the existing duck-typed precedent three lines above, 'if hasattr(samples, total_accepted_samples)', which adds Total Accepted Samples / Acceptance Ratio for MCMC searches — search-specific blocks are already the idiom. Guard with .get() so other searches are unaffected.

Naming constraint: emit these as neutral factual counts. Do NOT label them a smoothness metric in user-facing output — the resurrection-rate to HMC-divergence-rate correlation is unvalidated (that validation is a separate ideas.md item, wsdev#117 resurrection diagnostics).

Deliberately out of scope: making resurrect trigger on non-finite gradients. That would change search behaviour and shift every existing benchmark number, so the wsdev #117/#125 pix results would stop being comparable without re-running. Decide the resurrection policy AFTER the counters show how often frozen lanes actually occur.

Motivation: pixelized-mesh MultiStartProdigy campaigns (wsdev #117/#125) judge mesh differentiability indirectly, from final logL and raw resurrection counts. The value/gradient NaN split directly answers the open question in pix_prodigy_laptop_gpu_findings.md section 6.2 — whether DelaunayNN's 109 free-AdaptSplit lane deaths were NaN deaths like plain Delaunay or survivable over-regularized-floor deaths like knn — and would have caught frozen zombie lanes in every run to date, some of which may have been misattributed to regularization plateaus.

Sizing note: this is SMALL despite the prose length — roughly a counter plus two dict keys in search.py, and a ~5-line guarded block in text_util.py, plus a benchmark and unit tests.

<!-- formalised by the Intake (Conception) Agent on 2026-08-14 from user-intake -->

[Jammy2211]

Library PR Created

• PyAutoFit: feat: MultiStartGradient value/gradient NaN step diagnostics #1473 (pending-release)
• autolens_profiling: searches: standalone overhead demo for MultiStartGradient NaN accounting autolens_profiling#127 (pending-release) — the overhead demonstration; merge after feat: MultiStartGradient value/gradient NaN step diagnostics #1473, since its numbers describe the merged library code.

API impact

None — internal changes only. Two additive samples_info keys and gradient-search-only search.summary lines. grep across autofit_workspace, autogalaxy_workspace, autolens_workspace, autolens_workspace_test and HowToLens for the changed surface returns zero references, so no workspace migration is owed.

Smoke tests — all six workspaces

Run via pyauto-heart smoke --root <task worktree>, so the branch's PyAutoFit was the imported source (confirmed: Heart's runner replaces PYTHONPATH from --root and preflights that every module resolves under it).

workspace	result
autofit	pass
autogalaxy	pass
autolens	pass
autolens_test	3 failed (pre-existing — see below)
euclid	6/6 passed
HowToLens	50 passed, 0 failed

The three failures are pre-existing on main

interferometer/jax_likelihood/rectangular.py
interferometer/jax_likelihood/mge.py
multi_dataset/jax_likelihood/mge.py

All are stale pinned-value assertions, marginally over tolerance:

AssertionError: Not equal to tolerance rtol=0.0001, atol=0
 [0]: -3163.8939270532364 (ACTUAL), -3164.286252 (DESIRED)
Max relative difference among violations: 0.00012399

Control run: the same workspace, same harness, same cached environment, with --root pointed at canonical main produced the same three failures with byte-identical ACTUAL and DESIRED values. They are not caused by this branch — and independently, the failing scripts contain no MultiStart reference at all, so they never execute the changed code.

Filed as PyAutoMind/draft/bug/workspaces/jax_likelihood_pins_stale_by_1e4.md.

Verification beyond the test suite

• Device-side jnp.all(jnp.isfinite(grad)) cross-checked against host recomputation on real gradients straddling a non-differentiability cliff — exact agreement.
• Gradient-NaN counter proven to fire: forced a value-finite/gradient-NaN lane via the jnp.where AD trap → 14/80 lane-steps counted, n_resurrections unchanged (frozen lanes counted, deliberately not resurrected).
• search.summary verified on disk, with rates checked against count / (n_starts × total_steps) read off the file.
• Both guards mutation-tested — removing either fails tests, including the pre-existing test__search_summary_to_file.

Known gap

Resume accumulation is unit-test-only. Demonstrating it end-to-end requires resuming a killed mid-run search, which fails on main for unrelated reasons (FoM sanity check compares a stored log-likelihood against the multi-start chi-squared convention). Filed as PyAutoMind/draft/bug/autofit/multistart_gradient_resume_fom_sanity_check.md.

Heart was YELLOW at ship time (score 70, red_reasons: []); both reasons — pre-existing autogalaxy notebook validation failures and tenant-firewall manifest drift — are unrelated to this branch and were acknowledged by the human before shipping.

[Jammy2211]

Shipped

• PyAutoFit#1473 MERGED as fbfcece3 — verified present on origin/main.
• autolens_profiling#127 MERGED as a34d6191.

All checks green before merge (docs-build, unittest 3.12 + 3.13; profiling lint). Smoke clean for this branch across all six workspaces; the three autolens_test jax_likelihood failures reproduce byte-identically on main and are filed separately.

Shipped variant costs 4.1 us on a 1.03 s step (0.0004% of run time).

Two bugs spun off rather than absorbed into this PR:

• draft/bug/autofit/multistart_gradient_resume_fom_sanity_check.md — blocks end-to-end resume verification of these counters.
• draft/bug/workspaces/jax_likelihood_pins_stale_by_1e4.md

Still owed: the GPU row for the profiling artifact (unchecked box in autolens_profiling#127).