refactor: NSS onto Fitness, trace preflight and x64 check (search-extensibility A3b, stacked on A3)

[Jammy2211]

Closes #1678

Search-extensibility phase A3b. Stacked on A3 (#1680, itself stacked on A2 #1679): review and merge those first. The diff against feature/search-ext-a3-samples-checkpointer is the three commits below plus a merge of the updated A3 branch.

Summary

1. NSS onto Fitness and run(ctx) (aeadaae48)
   • NSS now samples through its fit's scalar objective (Fitness.objective via make_fitness, from its declared objective_target and invalid_value). It no longer uses the inline nss_log_likelihood_from closure, so NSS shares the one guard path: NaN/inf handling, the configured magnitude ceiling and the traced xp.where assertion penalty. An assertion no longer raises TracerBoolConversionError under jit.
   • NSS moves onto run(ctx) + raw_samples_from and resumes from ctx.resume (NativeFileCheckpointer("nss_checkpoint.pkl")).
   • The bridge builds the Fitness before sampling. The resume likelihood sanity check therefore runs at the start of the run.
   • make_fitness is now the only place under autofit/non_linear/search that builds a Fitness.
   • Removed nss_log_likelihood_from, which had no external users.
2. JAX trace preflight and x64 check (8bc269743, new autofit/non_linear/search/preflight.py). For a JAX analysis only, these run in start_resume_fit after the test-mode bypass and the REQUIRED gate.
   • check_x64 warns once per process when jax_enable_x64 is off. It raises for a search declaring the new requires_fp64 capability (default False; no search sets it).
   • trace_preflight runs one jax.eval_shape of the declared objective kind (batched or scalar) at the prior medians, plus the gradient kind when gradient == uses. Nothing is compiled or executed.
   • A likelihood that cannot be traced raises SearchException, which names the search and the kind and is chained to the tracer error.
   • PYAUTO_JAX_PREFLIGHT=0 skips the preflight. PYAUTO_JAX_PREFLIGHT_PROBE=1 adds an eager numerical probe, off by default.
   • docs/design/run_ctx.md gains Revision 2. The signature and members are unchanged.
3. Dynesty single-core control flow (c89cbc430)
   • AbstractDynesty._fit chose the single-core path by raising a RuntimeError inside a try. That meant any RuntimeError raised during a pooled run (including JAX's XlaRuntimeError) silently restarted the run single-core.
   • The path is now chosen by a private _SingleCoreRun sentinel. Only a RuntimeError from creating the pool still falls back, chained as the cause.

Tests run

• pytest test_autofit -q -x on the merged branch (local, numpy 2.2.6, JAX present): 3555 passed, 1 skipped, 4 xfailed.
• New and changed tests:
  • nss/test_nss_fitness.py: assertions under jit, plus an NSS fit with assertions under jit.
  • test_preflight.py:
    • an np.asarray likelihood raises with the error chained;
    • batched and gradient kinds;
    • the env-flag skip, the numpy no-op and the test-mode-2 bypass;
    • the REQUIRED gate runs before the preflight;
    • the probe is off by default and warns on an invalid value;
    • fp32 warns once, and raises when requires_fp64 is set.
  • nest/test_dynesty_single_core_sentinel.py.
  • Updated nss/test_checkpoint.py, nss/test_log_likelihood_ceiling.py and test_make_fitness.py.
• Not run here: the phase witness (NSS fitting gaussian_x3_blend with its ordered-centre assertions through the autofit_inference harness). The unit-scale analogue (test_nss_fit_with_assertions_runs_under_jit) passes.

Decisions taken

• New capability requires_fp64 (default False, no search sets it): added to the manifest in a schema-compatible way. fp32 only warns until a search opts in.
• Numerical probe off by default (PYAUTO_JAX_PREFLIGHT_PROBE=1 turns it on). The default preflight stays eval_shape only, with no likelihood evaluation, so it costs no compile.
• Preflight placement: after the test-mode bypass and after the REQUIRED gate, so the D2 regression (REQUIRED + PYAUTO_DISABLE_JAX=1 + PYAUTO_TEST_MODE=2 completes) is untouched.
• The bridge repeats the REQUIRED gate before building the Fitness, for callers that invoke _fit directly.
• nss_log_likelihood_from removed outright rather than deprecated, because it had no external users.
• Dynesty fallback is narrowed to a RuntimeError raised while creating the pool. A failure during the run now surfaces instead of silently restarting single-core.

🤖 Generated with Claude Code