feat: search capability declarations, fail-fast gate, lazy registry and run(ctx) design note (search-extensibility A1)

[Jammy2211]

Summary

Phase A1 of the search-extensibility epic (PyAutoFit#1674). Every search now declares its capabilities as class attributes; a declarative, lazy registry mirrors them with a completeness test and a versioned python -m autofit search-manifest --json (search-manifest@1); Analysis.is_jax replaces the nine scattered _use_jax probes and factor graphs derive their backend from their factors (D12); a JAX-required search given a numpy analysis raises one shared SearchException after the test-mode bypass (D2); Nautilus(force_x1_cpu=True) works with numpy; the search API page, an RTD capability matrix and a canonical citations page are generated from the manifest with a --check job (D18); and docs/design/run_ctx.md freezes the run(ctx) signature and FitContext members for A2 (D7). Eight commits in that order. Shipped under the human's 2026-10-08 --auto launch at effective level supervised; tier judge → human /prm. Merge this before the three downstream docs PRs (their intersphinx label resolves once RTD latest rebuilds).

Golden identifiers unchanged: no capability attribute is in any __identifier_fields__ (tested).

Decision taken (decide-and-flag, one)

Factor-graph backend rules. FactorGraphModel and ModelAnalysis now default use_jax=None and derive the flag from their factors (the plan's "a FactorGraphModel inherits use_jax"); an explicit use_jax=True with disagreeing factors raises; factors declaring different gradient_modes fall back to "reverse"; the agreement check sits next to the REQUIRED gate, after the test-mode bypass, so a mixed graph under PYAUTO_TEST_MODE>=2 is not caught. Rejected alternative: check before the bypass (stricter, riskier for smoke runs) and keep use_jax=False as the default (would not inherit). One-command revert of the placement: move the isinstance(analysis, FactorGraphModel) block in abstract_search.py above mode = test_mode_level(); the default change reverts with git revert 196abf39a (which also removes is_jax).

Judgement values for the reviewer (one line each in the class and registry.py): status experimental for NSS, SMC, ADABelief, Lion; warm_start provider for BFGS/LBFGS/MultiStart (the Brain samplers faculty's "mode-finders provide a point"); resumable False for Emcee (report §3.1 literally); batched True for NUTS (vmaps over chains). The §3.1 checkpointer attribute is deferred to A3.

API Changes

Added Analysis.is_jax; AnalysisFactor.is_jax/.gradient_mode; FactorGraphModel.factors_disagreeing_on_backend(), .check_backend_agreement(), .gradient_mode; ModelAnalysis.gradient_mode; the 15 capability class attributes on every search; modules autofit.non_linear.search.capabilities and autofit.non_linear.search.registry; the python -m autofit search-manifest CLI. Changed: FactorGraphModel(use_jax=None)/ModelAnalysis(use_jax=None) derive the flag (was False); HierarchicalFactor honours PYAUTO_DISABLE_JAX; JAX-required searches raise SearchException instead of ValueError; search.summary/model.info start with a capability header; search_summary_to_file(..., search_capabilities=None); FactorGraphModel.tree_flatten aux carries use_jax.
See full details below.

Test Plan

• pytest test_autofit -x: 3409 passed, 2 skipped, 9 xfailed (the A0 strict xfails only, none added)
• nojax emulation over non_linear, analysis and the generated-docs test: 1192 passed, 132 skipped, 7 xfailed; import autofit imports no optional backend (checked via sys.modules)
• Sphinx full env and emulated minimal env: 30 warnings, equal to docs/sphinx_warning_baseline.txt; matrix renders 15 rows; docs/_generate_searches.py --check up to date
• PyAutoLens 832 passed / 1 xfailed; PyAutoGalaxy 1357 passed
• afW searches/{mcmc,nest,mle} and afT Nautilus, NSS exit 0; afT BlackJAXNUTS/MultiStartAdam fail their own accuracy asserts identically on main (control run)
• REQUIRED + PYAUTO_DISABLE_JAX=1 PYAUTO_TEST_MODE=2 completes for NUTS, MultiStartAdam, NSS
• CI green on unittest 3.12 / 3.13 / nojax, docs, and the new generated-search-docs job

Full API Changes (for automation & release notes)

Removed

• none (the _use_jax probes are replaced internally; Analysis.use_jax is unchanged)

Added

• Analysis.is_jax — read-only, the single JAX probe
• AnalysisFactor.is_jax, AnalysisFactor.gradient_mode
• FactorGraphModel.factors_disagreeing_on_backend(), FactorGraphModel.check_backend_agreement(), FactorGraphModel.gradient_mode
• ModelAnalysis.gradient_mode
• Class attributes on every search: jax_use, gradient, batched, honours_gradient_mode, posterior_kind, produces_evidence, resumable, warm_start, install_extra, upstream_url, citation_keys, status, test_mode_budget, objective_target, invalid_value
• autofit.non_linear.search.capabilities: JaxUse, Gradient, PosteriorKind, WarmStart, Status, ObjectiveTarget, CAPABILITY_ATTRIBUTES, capabilities_from, JAX_REQUIRED_MESSAGE, check_jax_required, capability_summary_from
• autofit.non_linear.search.registry: SEARCHES, entries(), entry(), manifest(), MANIFEST_SCHEMA
• python -m autofit search-manifest [--json]
• docs/_generate_searches.py --check; docs/searches/{index,citations}.rst; docs/design/run_ctx.md

Migration

• Before: getattr(analysis, "_use_jax", False) → After: analysis.is_jax
• Before: a JAX-required search with a numpy analysis raised ValueError → After: autofit.exc.SearchException (message JAX_REQUIRED_MESSAGE)
• Before: FactorGraphModel(use_jax=False) default → After: use_jax=None derives from the factors. Explicit FactorGraphModel(use_jax=False) does not force JAX factors onto numpy: when it is fitted as one analysis, the agreement check rejects any factor that disagrees and raises SearchException. To fit the whole graph on numpy, build every factor's analysis (and any HierarchicalFactor) with use_jax=False and leave FactorGraphModel(use_jax=...) unset

Validation checklist (--auto run — plan was not pre-approved)

• Effective level: supervised (header: supervised, cap: feature@large → supervised); ship checkpoint resolved by decide-and-flag (one decision, above)
• Plan: on the issue (feat: search capability declarations, fail-fast gate, lazy registry and run(ctx) design note (search-extensibility A1) #1674), written at start, unmodified since
• Gate: tests 3409 pass / 2 skip / 9 xfail + nojax 1192 pass + downstream PyAutoLens 832 / PyAutoGalaxy 1357 · smoke afW searches ×3 + afT Nautilus/NSS exit 0 (NUTS/MultiStartAdam pre-existing accuracy failures on main) · review CLEAN (Fable session over the Opus-written branch; witness basis-cited: manifest 15 rows, sys.modules free of optional backends, SearchException raised for NUTS+numpy, generator --check current, design note present, 54 registry/docs tests pass) · Heart STALE release validation incomplete: no rehearsal for current source
• Human: plan sound in hindsight?
• Human: diff matches plan (no scope creep)? Check the flagged decision and the judgement values
• Human: merge, amend, or reject — then log the outcome

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

Independent adversary review (Codex gpt-6-astra), enacted 2026-10-08

Review archived at PyAutoMind draft/research/autofit/search_extensibility_epic_reviews/04_codex_astra_a1_pr1675.md. Witness: all five clauses held (each of the 15 registry deletions detected; manifest 15 rows; shared SearchException for all seven REQUIRED searches; import free of optional backends; design note present). Verdict FINDINGS (6); all six reproduced and enacted in follow-up commits 2c07ebea6..53e39c9c0:

1. P1 legacy FactorGraphModel pickles with _use_jax=False came back is_jax=True and could not flatten → __setstate__ migration + test · 2. P1 a graph over a HierarchicalFactor(use_jax=True) inferred is_jax=False → backend derived from the flattened factors (the same set the agreement check reads) + test · 3. ModelAnalysis(graph) bypassed the whole-graph check → the gate unwraps ModelAnalysis/AnalysisFactor first + test · 4. NSS objective_target.space is physical (draws are transformed before algo.init) · 5. BFGS/LBFGS/MultiStart invalid_value is what the backend sees on NaN, +inf (×−2 of the −inf replacement), with the conditional −inf on FitException documented and pinned by a test for A3's sentinel normalisation · 6. Emcee resumable=True (HDF backend resume), correcting the report §3.1 table.
   Kept on the adversary's recommendation: gates after the test-mode bypass, the use_jax=None inheritance intent, NUTS batched=True, the experimental statuses and mode-finder warm_start=provider. After the fixes: 3417 passed / 2 skipped / 9 xfailed; golden table untouched; generator --check current; Sphinx 30 warnings = baseline; import autofit loads no optional backend. Nojax emulation: 1466 passed with 4 pre-existing order-dependent failures in graphical/ that reproduce identically on origin/main.