Warm-start abstraction for gradient MCMC: multi-chain vmapped BlackJAXNUTS + metric/start-point injection from a Result

[Jammy2211]

Phase 6 pre-req of the autolens_profiling inference programme (results/notes/inference/PROGRAMME.md §4 Phase 6; Gate C reworded 2026-08-24 to batched-pipeline value).

Scope (library, no GPU needed)

• af.BlackJAXNUTS: multi-chain via vmap (n_chains), inverse-mass-matrix injection (diagonal / dense / low-rank), start-point injection from a previous Result without touching priors — wire InitializerParamStartPoints + a covariance carrier.
• A warm-start abstraction: Result → (start points, metric); promote the wsdev _warm_start cache pattern (MLE + Laplace covariance) into PyAutoFit. Scientific priors stay untouched — start points and metric only.
• Diagnostics surfaced from blackjax 1.6: rank-normalised ESS_bulk/ESS_tail, split-R̂, divergence count + location, tree-depth distribution.
• Unit tests without JAX where possible (CPU-only, tiny models); no JAX in unit tests per repo convention.

Design constraints: ChEES-HMC / MAMS arms may follow — keep the chain-batching and metric plumbing sampler-agnostic.

[Jammy2211]

Design decision (2026-08-24, human): no new af.WarmStart class. The chaining API (result.model_centred*) does prior passing — it rewrites priors — which is exactly what a warm start must not do. Start points reuse the existing initializer API: InitializerParamStartPoints.from_result(result, model=None, point="max_log_likelihood"|"median_pdf", n_points=1, jitter=0.0, seed=0), path-matched onto the (possibly re-composed) model, priors untouched. The metric goes in via BlackJAXNUTS(inverse_mass_matrix=None|"diagonal"|"dense"|ndarray|Result|Samples) (a Result/Samples pulls samples.covariance_matrix; MLE-only results raise unless an explicit Laplace array is passed), plus num_chains, mass_matrix_shrinkage, share_adaptation. Chain-batching/metric helpers live in a sampler-agnostic blackjax/chains.py for later ChEES/MAMS arms.

nuts = af.BlackJAXNUTS(
    num_chains=16,
    initializer=af.InitializerParamStartPoints.from_result(result, n_points=16, jitter=0.1),
    inverse_mass_matrix=result,
)

Implementation in progress on feature/warm-start-multichain-nuts.

[Jammy2211]

Shipped: PR #1522 merged. InitializerParamStartPoints.from_result, Result.start_point / start_point_from(...), multi-chain vmapped BlackJAXNUTS with inverse_mass_matrix (None/diagonal/dense/ndarray/Result/Samples), mass_matrix_shrinkage, share_adaptation, and rank-normalised ESS / split-R̂ / divergence / tree-depth diagnostics on SamplesMCMC. Pre-merge smoke: 2 warm-started chains on a 2-param Gaussian, dense metric — R̂ 1.002/1.000, ESS_bulk ~1700 of 2000, 0 divergences. Follow-ups: low-rank metric, ChEES/MAMS arms (Phase 6 of the inference programme).