feat: MultiStartGradient value/gradient NaN step diagnostics

[Jammy2211]

Summary

MultiStartGradient detected a dead lane only via the value (alive = np.isfinite(np.asarray(foms))). The gradient was never checked inside the fit loop, so a lane whose likelihood is finite but whose gradient is non-finite was counted alive, had its update zeroed by optax.apply_if_finite, and silently froze in place — a differentiability failure that looks exactly like convergence in the figure-of-merit trace.

This counts both failure modes per step, disjointly, and surfaces them:

• value-NaN lane-steps — the likelihood is undefined (today's resurrection trigger)
• gradient-NaN lane-steps — the likelihood is defined but not differentiable (new; previously unmeasured)

Gradient finiteness is reduced inside the jitted call, as a third output. Reducing eagerly outside the jit measured ~+3%/step against a ~1.5% noise floor — worse than pulling the whole gradient to host, which it was meant to avoid — because an un-jitted reduction costs a kernel dispatch plus a host round-trip. On a real MGE lens likelihood the shipped variant costs 4.1 µs on a 1.03 s step = 0.0004% of run time (measured in PyAutoLabs/autolens_profiling#PENDING).

Measurement only. resurrect still triggers on value-NaN alone, so every existing benchmark number stays comparable. The resurrection policy is decided after the counters show how often frozen lanes actually occur.

Closes #1472

API Changes

None — internal changes only.

Two additive samples_info keys (n_value_nan_lane_steps, n_grad_nan_lane_steps) and extra search.summary lines for gradient searches. No public symbol added, removed, renamed, or re-signatured. search_summary_from_samples is on every search's summary path, so the new block is guarded on the key — searches without gradients (Nautilus) emit byte-identical output to before.

See full details below.

Test Plan

• Full suite green: 1747 passed, 2 skipped
• _nan_lane_counts unit-tested (pure NumPy, no JAX in the library suite), including both-NaN-on-the-same-lane disjointness and integer-dtype grad_finite
• samples_info plumbing tested via hand-built search_internal, including a legacy dict with neither key → 0
• End-to-end JAX fit across all four path combinations (batch_size None/2 × resurrect off/on); with resurrect=True the value-NaN count equals n_resurrections exactly
• 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, confirming frozen lanes are counted but deliberately not resurrected
• Device reduction verified against host recomputation on real gradients straddling a non-differentiability cliff — exact agreement
• search.summary verified on disk (not just in-process), with rates checked against count / (n_starts × total_steps) read off the file
• Both guards mutation-tested: dropping alive & fails exactly the same-lane cases; ungating the summary block fails the Nautilus test and the pre-existing test__search_summary_to_file

Known coverage gap

Resume accumulation is covered only by unit tests over hand-built search_internal dicts. Demonstrating it end-to-end requires resuming a killed mid-run search, which fails on main for unrelated reasons: the FoM sanity check compares a stored log-likelihood against the multi-start chi-squared convention (a consistent −2× relationship). Reproduced on main with the stock af.ex.Analysis, so it predates this PR. Filed as PyAutoMind/draft/bug/autofit/multistart_gradient_resume_fom_sanity_check.md; re-check the counters here once it is fixed.

Full API Changes (for automation & release notes)

Removed

• None

Added

• samples_info["n_value_nan_lane_steps"] — lane-steps where the likelihood was undefined
• samples_info["n_grad_nan_lane_steps"] — lane-steps where the likelihood was defined but not differentiable
• search_internal["n_value_nan_lane_steps"] / search_internal["n_grad_nan_lane_steps"] — persisted, resume-safe via .get(..., 0)

Changed Behaviour

• search.summary gains Resurrections, Value-NaN Lane-Steps, Gradient-NaN Lane-Steps and two normalised rates for gradient searches only. Guarded on the samples_info key, so all other searches are unaffected.
• MultiStartGradient's internal jitted objective returns a third output (per-start gradient finiteness). Internal to _fit; the single-point objective used for the initial draw is deliberately left as a 2-tuple.

Migration

• None required. Counters default to 0 when absent, so pre-existing search_internal files load unchanged.

Generated by the PyAutoLabs agent workflow.