From 5402c761956a36b796a1a257b500fe3888176a1f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 14 Oct 2026 09:30:37 +0100 Subject: [PATCH 1/6] fix: keep eager objective when unpickling a pre-A2 Fitness A Fitness pickled before the objective factory carries use_jax_vmap / use_jax_jit. __setstate__ discarded use_jax_jit and defaulted compile=True, so an eager (use_jax_jit=False) checkpoint resumed jitted and broke likelihoods valid only eagerly. Map the legacy flags onto compile exactly. Co-Authored-By: Claude Fable 5.1 --- autofit/non_linear/fitness.py | 10 +-- .../non_linear/test_fitness_objective.py | 66 +++++++++++++++++++ 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/autofit/non_linear/fitness.py b/autofit/non_linear/fitness.py index 8b8cfc313..f7631ee1a 100644 --- a/autofit/non_linear/fitness.py +++ b/autofit/non_linear/fitness.py @@ -960,12 +960,14 @@ def __setstate__(self, state): self.__dict__.setdefault("gradient_mode", None) # `Fitness` objects pickled before the objective factory (search-extensibility A2) # carry `use_jax_vmap` / `use_jax_jit` instead of `batched` / `compile`. Their - # dispatch maps onto the factory: vmap -> the batched objective; jit or not, a JAX - # scalar objective is now jitted lazily. + # dispatch is preserved exactly: vmap -> the (jitted) batched objective; jit -> the + # jitted scalar objective; neither -> the eager scalar objective. A pickle that + # predates both flags dispatched eagerly too, so it stays eager. legacy_vmap = self.__dict__.pop("use_jax_vmap", None) - self.__dict__.pop("use_jax_jit", None) + legacy_jit = self.__dict__.pop("use_jax_jit", None) + if "compile" not in self.__dict__: + self.__dict__["compile"] = bool(legacy_vmap) or bool(legacy_jit) self.__dict__.setdefault("batched", bool(legacy_vmap)) - self.__dict__.setdefault("compile", True) self._objectives = {} self._call = self.objective(BATCHED if self.batched else SCALAR) diff --git a/test_autofit/non_linear/test_fitness_objective.py b/test_autofit/non_linear/test_fitness_objective.py index d01abca41..b32d7cb91 100644 --- a/test_autofit/non_linear/test_fitness_objective.py +++ b/test_autofit/non_linear/test_fitness_objective.py @@ -225,3 +225,69 @@ def test_evaluate_in_chunks_matches_one_call_and_compiles_one_shape(): assert np.allclose(np.asarray(chunked), np.asarray(whole)) # Shapes (2, d) and (5, d): the padded ragged chunk reuses the (2, d) executable. assert batched.__wrapped__._cache_size() == 2 + + +class EagerOnlyAnalysis(af.ex.Analysis): + """ + A JAX analysis whose likelihood is valid eagerly but not under `jax.jit`: it + converts a parameter to a Python float, which raises on a tracer. + """ + + def log_likelihood_function(self, instance, **kwargs): + return float(np.asarray(instance.centre)) * 0.0 + 50.0 + + +def _legacy_state(fitness, use_jax_vmap, use_jax_jit): + """ + The `__dict__` a pre-A2 (A1-era) `Fitness` pickled: `use_jax_vmap` / `use_jax_jit` + instead of `batched` / `compile`. + """ + state = fitness.__getstate__() + state.pop("batched") + state.pop("compile") + state["use_jax_vmap"] = use_jax_vmap + state["use_jax_jit"] = use_jax_jit + return state + + +def _restore(state): + import pickle + + restored = Fitness.__new__(Fitness) + restored.__setstate__(state) + return pickle.loads(pickle.dumps(restored)) + + +def test_legacy_eager_pickle_stays_eager(): + pytest.importorskip("jax") + + model = af.Model(af.ex.Gaussian) + analysis = EagerOnlyAnalysis( + data=np.ones(20), noise_map=np.ones(20) * 0.1, use_jax=True + ) + fitness = Fitness(model=model, analysis=analysis, compile=False) + vector = _vector(fitness) + assert float(fitness._call(vector)) == pytest.approx(50.0) + + restored = _restore(_legacy_state(fitness, use_jax_vmap=False, use_jax_jit=False)) + + assert restored.compile is False + assert restored.batched is False + assert not hasattr(restored._call, "__wrapped__") + assert float(restored._call(vector)) == pytest.approx(50.0) + + +def test_legacy_jit_and_vmap_pickles_stay_compiled(): + pytest.importorskip("jax") + + fitness = _fitness(use_jax=True) + + jitted = _restore(_legacy_state(fitness, use_jax_vmap=False, use_jax_jit=True)) + assert jitted.compile is True + assert jitted.batched is False + assert hasattr(jitted._call, "__wrapped__") + + vmapped = _restore(_legacy_state(fitness, use_jax_vmap=True, use_jax_jit=False)) + assert vmapped.compile is True + assert vmapped.batched is True + assert hasattr(vmapped._call, "__wrapped__") From 1f33755f5869ec91ec60780423edf8c67304936c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 14 Oct 2026 09:36:32 +0100 Subject: [PATCH 2/6] fix: preflight validates the resolved execution settings the backend runs Preflight traced the class-declared objective kind through a throwaway Fitness that ignored runtime options, so it rejected an eager use_jax_jit=False likelihood and a MultiStartAdam gradient_mode='reverse' override. NonLinearSearch.execution_settings(analysis) now resolves batching, compile, gradient use and gradient mode once (preflight.ExecutionSettings); preflight traces exactly that and the run(ctx) bridge, Dynesty and MultiStartGradient build their Fitness from the same object. Co-Authored-By: Claude Fable 5.1 --- autofit/non_linear/search/abstract_search.py | 18 ++- .../search/mle/multi_start_gradient/search.py | 9 +- .../search/nest/dynesty/search/abstract.py | 13 +- autofit/non_linear/search/preflight.py | 138 +++++++++++++++--- .../non_linear/search/test_preflight.py | 82 +++++++++++ 5 files changed, 232 insertions(+), 28 deletions(-) diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index b5528b393..c067910b8 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -1476,7 +1476,9 @@ def _fit(self, model: AbstractPriorModel, analysis: Analysis): cap.check_jax_required(self, analysis) fitness = self.make_fitness( - analysis=analysis, model=model, **self.fitness_overrides(analysis) + analysis=analysis, + model=model, + **self.execution_settings(analysis).fitness_overrides, ) ctx = FitContext( @@ -1523,6 +1525,20 @@ def fitness_overrides(self, analysis: Analysis) -> dict: """ return {} + def execution_settings(self, analysis: Analysis): + """ + How this search executes its objective on ``analysis``, resolved once: batching, + eager or jitted execution, gradient use and the resolved gradient mode, plus the + ``Fitness`` keyword arguments the backend builds its ``Fitness`` with + (``autofit.non_linear.search.preflight.ExecutionSettings``). + + The JAX trace preflight validates exactly this object and the backend builds its + ``Fitness`` from it, so the two cannot disagree. + """ + from autofit.non_linear.search.preflight import ExecutionSettings + + return ExecutionSettings.resolve(self, analysis) + def make_fitness( self, analysis: Analysis, model: AbstractPriorModel, **overrides ) -> Fitness: diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 490036fb9..044fad1b2 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -1003,10 +1003,15 @@ def _fit( else self.iterations_per_quick_update ) + settings = self.execution_settings(analysis) + fitness = self.make_fitness( analysis=analysis, model=model, - iterations_per_quick_update=iterations_per_quick_update, + **{ + **settings.fitness_overrides, + "iterations_per_quick_update": iterations_per_quick_update, + }, ) # -2 * log_posterior, to MINIMIZE. value_and_grad over every start. @@ -1031,7 +1036,7 @@ def _fit( # analysis's declaration and this search's override, and used for both # objectives below. Forward and reverse give the same ``(value, grad)``; # everything downstream of the two builders is mode-agnostic. - gradient_mode = self._resolved_gradient_mode(analysis) + gradient_mode = settings.gradient_mode if not self.silence: source = ( "search override" diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index f69232412..9a42f3d31 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -211,6 +211,14 @@ def __init__( self.logger.debug(f"Creating {self.__class__.__name__} Search") + def fitness_overrides(self, analysis) -> dict: + """ + A JAX analysis is jitted lazily by the scalar objective `__call__` dispatches to; + the deprecated `use_jax_jit=False` keeps its old meaning (eager JAX) for one + release. + """ + return {"compile": getattr(self, "use_jax_jit", None) is not False} + @property def search_kwargs(self) -> Dict: """Shared search kwargs passed to both Static and Dynamic Dynesty samplers.""" @@ -271,10 +279,7 @@ def _fit( fitness = self.make_fitness( analysis=analysis, model=model, - # A JAX analysis is jitted lazily by the scalar objective `__call__` - # dispatches to; the deprecated `use_jax_jit=False` keeps its old meaning - # (eager JAX) for one release. - compile=self.use_jax_jit is not False, + **self.execution_settings(analysis).fitness_overrides, ) if not isinstance(self.paths, NullPaths): diff --git a/autofit/non_linear/search/preflight.py b/autofit/non_linear/search/preflight.py index 3b31e0b4d..ff10ed542 100644 --- a/autofit/non_linear/search/preflight.py +++ b/autofit/non_linear/search/preflight.py @@ -11,9 +11,13 @@ environment variable read at ``autonerves`` import, so importing ``jax`` first runs every likelihood in fp32 without a word. Warn once per process; raise for a search that declares ``requires_fp64 = True`` (none does today). -- ``trace_preflight``: one ``jax.eval_shape`` of the search's declared objective kind - (``"batched"`` for a ``batched`` search, else ``"scalar"``) at the prior-median vector, - plus one of its gradient kind when the search declares ``gradient == "uses"``. That is +- ``trace_preflight``: one ``jax.eval_shape`` of the objective kind the search actually + runs (``"batched"`` for a batched execution, else ``"scalar"``) at the prior-median + vector, plus one of its gradient kind (in the gradient mode the search resolves) when + the search declares ``gradient == "uses"``. The kinds and modes come from the search's + ``ExecutionSettings`` (``NonLinearSearch.execution_settings``), the same resolved object + the backend builds its ``Fitness`` from, so an eager (``compile=False``) scalar + execution is not traced at all and a search's ``gradient_mode`` override is honoured. That is one trace and no compile or execution, so it costs seconds at most where the backend's own first call would compile for minutes, and a non-traceable likelihood (an ``np.asarray`` or ``float()`` on a parameter-dependent value, a Python ``if`` on an @@ -36,6 +40,8 @@ import logging import os import warnings +from dataclasses import dataclass, field +from typing import Optional from autofit import exc from autofit.non_linear.objective import ( @@ -131,41 +137,127 @@ def check_x64(search, analysis): warnings.warn(X64_WARNING, UserWarning, stacklevel=3) -def preflight_kinds(search) -> tuple: +@dataclass(frozen=True) +class ExecutionSettings: """ - The objective kinds the trace preflight traces for ``search``: its declared - execution kind, plus the matching gradient kind when it declares - ``gradient == "uses"``. + How a search will execute its objective on one analysis, resolved once from the + search's capability declarations, its runtime options and the analysis: the one + object the trace preflight validates and the backend builds its ``Fitness`` from + (``NonLinearSearch.execution_settings``). + + Attributes + ---------- + batched + Whether the backend evaluates the batched (``jax.vmap``) objective. A search may + resolve this at runtime (``Nautilus(use_jax_vmap=False)`` runs the scalar one). + compile + Whether the scalar objective is jitted. ``False`` is eager execution (the + deprecated ``use_jax_jit=False``), which never traces the scalar likelihood. + uses_gradient + Whether the backend differentiates the objective (``gradient == "uses"``). + gradient_mode + The resolved AD mode: the search's ``gradient_mode`` override when it honours + one, else the analysis's declaration (``autofit.jax.gradient``). + fitness_overrides + The ``Fitness`` keyword arguments the backend passes to ``make_fitness``, + including the resolved ``compile`` / ``batched`` / ``gradient_mode`` where they + differ from the defaults. """ + + batched: bool + compile: bool + uses_gradient: bool + gradient_mode: str + fitness_overrides: dict = field(default_factory=dict) + + @property + def kinds(self) -> tuple: + """ + The objective kinds the backend evaluates: its execution kind, plus the matching + gradient kind when it differentiates. + """ + kinds = [BATCHED if self.batched else SCALAR] + if self.uses_gradient: + kinds.append(BATCHED_VALUE_AND_GRAD if self.batched else VALUE_AND_GRAD) + return tuple(kinds) + + @classmethod + def resolve(cls, search, analysis) -> "ExecutionSettings": + """ + Resolve ``search``'s execution on ``analysis``: its ``fitness_overrides`` + (runtime batching, eager execution, batch size, ...) over its class + declarations, and its ``gradient_mode`` override (when the search honours one) + over the analysis's declared mode. + """ + from autofit.jax.gradient import resolve_gradient_mode + + search_cls = type(search) + overrides = dict(search.fitness_overrides(analysis)) + + override_mode: Optional[str] = overrides.get("gradient_mode") + if override_mode is None and search_cls.honours_gradient_mode: + override_mode = getattr(search, "gradient_mode", None) + if override_mode is not None: + overrides["gradient_mode"] = override_mode + + return cls( + batched=bool(overrides.get("batched", search_cls.batched)), + compile=bool(overrides.get("compile", True)), + uses_gradient=search_cls.gradient == cap.Gradient.USES, + gradient_mode=resolve_gradient_mode(analysis, override=override_mode), + fitness_overrides=overrides, + ) + + +def preflight_kinds(search, analysis=None) -> tuple: + """ + The objective kinds the trace preflight traces for ``search``: the kinds of its + resolved ``ExecutionSettings`` on ``analysis``, or of its class declarations when + no analysis is given. + """ + if analysis is not None: + return search.execution_settings(analysis).kinds + cls = type(search) - kinds = [BATCHED if cls.batched else SCALAR] - if cls.gradient == cap.Gradient.USES: - kinds.append(BATCHED_VALUE_AND_GRAD if cls.batched else VALUE_AND_GRAD) - return tuple(kinds) + return ExecutionSettings( + batched=bool(cls.batched), + compile=True, + uses_gradient=cls.gradient == cap.Gradient.USES, + gradient_mode="reverse", + ).kinds -def _preflight_fitness(search, analysis, model): +def _preflight_fitness(search, analysis, model, settings: ExecutionSettings): """ - A throwaway ``Fitness`` in the search's declared convention, without paths (so no - resume sanity check evaluates the likelihood) and without quick updates. + A throwaway ``Fitness`` built from the search's resolved execution settings, without + paths (so no resume sanity check evaluates the likelihood) and without quick + updates. """ return search.make_fitness( analysis=analysis, model=model, - paths=None, - iterations_per_quick_update=None, - background_quick_update=False, - live_visual_update=False, + **{ + **settings.fitness_overrides, + "paths": None, + "iterations_per_quick_update": None, + "background_quick_update": False, + "live_visual_update": False, + }, ) def trace_preflight(search, analysis, model): """ - Trace the search's declared objective kind(s) once with ``jax.eval_shape`` at the + Trace the objective kind(s) the search resolves for ``analysis`` + (``NonLinearSearch.execution_settings``) once with ``jax.eval_shape`` at the prior-median vector, raising a ``SearchException`` that names the search and chains the tracer error when the likelihood cannot be traced. Then run the numerical probe if ``PYAUTO_JAX_PREFLIGHT_PROBE`` asks for it. + An eager scalar execution (``compile=False``) is never traced by the backend, so its + scalar kind is not traced here either; a gradient kind is, since differentiation + traces even an eager objective. + A no-op for a numpy analysis and when ``PYAUTO_JAX_PREFLIGHT=0``. """ if not analysis.is_jax or not _env_flag(PREFLIGHT_ENV, default=True): @@ -174,11 +266,15 @@ def trace_preflight(search, analysis, model): import jax import jax.numpy as jnp - fitness = _preflight_fitness(search, analysis, model) + settings = search.execution_settings(analysis) + + fitness = _preflight_fitness(search, analysis, model, settings) vector = jnp.asarray(model.physical_values_from_prior_medians) - for kind in preflight_kinds(search): + for kind in settings.kinds: + if kind == SCALAR and not settings.compile: + continue objective = fitness.objective(kind, compile=False) argument = vector[None, :] if kind in (BATCHED, BATCHED_VALUE_AND_GRAD) else vector try: diff --git a/test_autofit/non_linear/search/test_preflight.py b/test_autofit/non_linear/search/test_preflight.py index 08c98495c..e3c14f6a0 100644 --- a/test_autofit/non_linear/search/test_preflight.py +++ b/test_autofit/non_linear/search/test_preflight.py @@ -297,3 +297,85 @@ def test_no_public_search_requires_fp64(): from autofit.non_linear.search.registry import entries assert not any(entry.capabilities["requires_fp64"] for entry in entries()) + + +# Review 05 finding 3: the preflight validates the *resolved* execution settings the +# backend runs (eager execution, the search's gradient override, its actual batching), +# not the class declarations. + + +def test_preflight_accepts_an_eager_likelihood_when_the_search_runs_eagerly( + monkeypatch, +): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + search = af.DynestyStatic(use_jax_jit=False) + + settings = search.execution_settings(NumpyCastAnalysis()) + assert settings.compile is False + assert settings.fitness_overrides["compile"] is False + + # Eager scalar execution never traces, so the preflight has nothing to reject. + preflight.trace_preflight(search, NumpyCastAnalysis(), _model()) + + reached = [] + monkeypatch.setattr( + search, "_fit", lambda *args, **kwargs: reached.append(True) or 1 / 0 + ) + with pytest.raises(ZeroDivisionError): + search.fit(model=_model(), analysis=NumpyCastAnalysis()) + assert reached == [True] + + +def _custom_vjp_analysis(declared_mode): + @jax.custom_vjp + def reverse_only(x): + return x + + def _fwd(x): + return x, None + + def _bwd(_, g): + return (g,) + + reverse_only.defvjp(_fwd, _bwd) + + class ReverseOnlyAnalysis(af.Analysis): + gradient_mode = declared_mode + + def __init__(self): + super().__init__(use_jax=True) + + def log_likelihood_function(self, instance): + return -0.5 * reverse_only(instance.centre - 50.0) ** 2 + + return ReverseOnlyAnalysis() + + +def test_preflight_honours_the_search_gradient_mode_override(): + analysis = _custom_vjp_analysis("forward") + + # The analysis declares forward mode, which a custom_vjp function cannot run ... + with pytest.raises(exc.SearchException, match="value_and_grad"): + preflight.trace_preflight(af.MultiStartAdam(), analysis, _model()) + + # ... but the search overrides it with reverse mode, which is what the backend runs. + search = af.MultiStartAdam(gradient_mode="reverse") + assert search.execution_settings(analysis).gradient_mode == "reverse" + preflight.trace_preflight(search, analysis, _model()) + + +def test_preflight_follows_the_resolved_batching(): + search = af.Nautilus(use_jax_vmap=False) + settings = search.execution_settings(TraceableAnalysis()) + assert settings.batched is False + assert settings.kinds == ("scalar",) + + with pytest.raises(exc.SearchException, match="'scalar'"): + preflight.trace_preflight(search, NumpyCastAnalysis(), _model()) + + +def test_preflight_still_rejects_an_untraceable_likelihood_of_a_jitted_search(): + with pytest.raises(exc.SearchException, match="TracerArrayConversionError"): + preflight.trace_preflight( + af.DynestyStatic(use_jax_jit=True), NumpyCastAnalysis(), _model() + ) From edcf351162f7366d9c5534f738e066a679505955 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 14 Oct 2026 09:44:59 +0100 Subject: [PATCH 3/6] fix: finalize the archive before writing .completed; NSS keeps its checkpoint start_resume_fit wrote .completed straight after save_results, before post_fit_output archived search_internal, so an archive write that failed (disk full) left a 'completed' fit the rerun loaded without ever retrying the archive; NSS also deleted its resume checkpoint inside run(). The fit now finalizes its retained state (archive, then resume-state discard) in NonLinearSearch._complete_fit before writing .completed, so a failure leaves the fit resumable and the rerun finalizes it. NSS no longer discards its checkpoint in run(). Co-Authored-By: Claude Fable 5.1 --- autofit/non_linear/search/abstract_search.py | 53 +++++++++-- autofit/non_linear/search/nest/nss/search.py | 16 +--- .../search/nest/nss/test_checkpoint.py | 90 +++++++++++++++++++ test_autofit/non_linear/test_checkpoint.py | 51 +++++++++++ 4 files changed, 189 insertions(+), 21 deletions(-) diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index c067910b8..967291fc4 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -744,7 +744,11 @@ class represented by model M and gives a score for their fitness. info=info, ) - if not self.paths.is_complete: + # `start_resume_fit` finalizes the archive itself, before it writes + # `.completed`; a completed fit's `post_fit_output` still tidies the folder. + resumed_or_started = not self.paths.is_complete + + if resumed_or_started: result = self.start_resume_fit( analysis=analysis, model=model, @@ -762,6 +766,7 @@ class represented by model M and gives a score for their fitness. self.post_fit_output( search_internal=result.search_internal, + finalize_search_internal=not resumed_or_started, ) gc.collect() @@ -963,10 +968,25 @@ def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Res model.unfreeze() - self.paths.completed() + self._complete_fit(search_internal=search_internal) return result + def _complete_fit(self, search_internal): + """ + Finalize the fit's retained internal state, then mark the fit complete. + + The order is the transaction: ``.completed`` is written only once the archive + (and the discard of the resume state it supersedes) has succeeded. An archive + write that fails (a full disk) therefore leaves no marker, the resume state is + still on disk, and the next run resumes and finalizes again, rather than + loading a "completed" fit whose archive never landed. + """ + if not skip_fit_output(): + self.finalize_search_internal(search_internal=search_internal) + + self.paths.completed() + def _test_mode_valid_parameter_vector( self, model: AbstractPriorModel, @@ -1192,7 +1212,21 @@ def result_via_completed_fit( return result - def post_fit_output(self, search_internal): + def finalize_search_internal(self, search_internal): + """ + Archive ``search_internal`` (``output_search_internal``) when the search retains + its internal state, otherwise remove the ``search_internal`` folder. + + Run before ``.completed`` is written (``_complete_fit``), so a failure here + leaves the fit resumable. + """ + if not self.retains_search_internal: + self.logger.info("Removing search internal folder.") + self.paths.remove_search_internal() + elif search_internal is not None: + self.output_search_internal(search_internal=search_internal) + + def post_fit_output(self, search_internal, finalize_search_internal: bool = True): """ Cleans up the output folderds after a completed non-linear search. @@ -1207,12 +1241,13 @@ def post_fit_output(self, search_internal): ---------- search_internal The internal search. + finalize_search_internal + Whether to archive (or remove) the internal state here. ``False`` when the + fit just ran, since ``start_resume_fit`` finalized it before writing + ``.completed``. """ - if not self.retains_search_internal: - self.logger.info("Removing search internal folder.") - self.paths.remove_search_internal() - elif search_internal is not None: - self.output_search_internal(search_internal=search_internal) + if finalize_search_internal: + self.finalize_search_internal(search_internal=search_internal) if not self.disable_output: self.logger.info("Removing all files except for .zip file") @@ -1346,7 +1381,7 @@ def _fit_bypass_test_mode( # Mark the fit complete, exactly as start_resume_fit does — a bypassed # fit must be resumable (paths.is_complete -> result_via_completed_fit # on the next run), or every rerun re-bypasses the whole pipeline. - self.paths.completed() + self._complete_fit(search_internal=None) return result diff --git a/autofit/non_linear/search/nest/nss/search.py b/autofit/non_linear/search/nest/nss/search.py index e2b5f5c3b..368f3a9b2 100644 --- a/autofit/non_linear/search/nest/nss/search.py +++ b/autofit/non_linear/search/nest/nss/search.py @@ -536,18 +536,10 @@ def _safe_ess(log_w_mean): seed=int(self.seed), ) - if checkpoint_path is not None and checkpoint_path.exists(): - try: - self.resume_state.discard(ctx.paths) - except OSError as exc: - self.logger.warning( - "Failed to delete completed-run checkpoint %s: %s. The " - "next fresh af.NSS fit at this path will attempt to resume " - "from it — delete manually if that is not desired.", - checkpoint_path, - exc, - ) - + # The checkpoint is kept until the result exists: `start_resume_fit` discards + # this resume state only after the archive is finalized and before it writes + # `.completed` (`NonLinearSearch._complete_fit`), so a failure while the result + # is processed leaves the run resumable. return search_internal @property diff --git a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py index e3789e3dd..2a9b34dca 100644 --- a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py +++ b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py @@ -236,3 +236,93 @@ def test__resume_state_loaded_when_ctx_resume_is_set(tmp_path): resume.load.assert_called_once_with(paths) assert isinstance(internal, nss_search_module._NSSInternal) + + +@requires_nss +def test__run_keeps_its_checkpoint_until_the_result_is_written(tmp_path): + """Review 05 finding 2: ``run(ctx)`` must not delete ``nss_checkpoint.pkl``. + + The resume state is discarded by ``NonLinearSearch._complete_fit`` only after the + archive is finalized, so a failure while the result is processed (after ``run`` + returns) leaves the run resumable. The run below resumes and terminates at once, + exactly as ``test__resume_state_loaded_when_ctx_resume_is_set``. + """ + from types import SimpleNamespace + from unittest.mock import MagicMock + + from autofit.non_linear.paths.null import NullPaths + + sentinel_state = SimpleNamespace( + integrator=SimpleNamespace(logZ=0.0, logZ_live=-100.0), + particles=SimpleNamespace( + position=jnp.zeros((4, 2)), + loglikelihood=jnp.zeros(4), + ), + ) + + # Minimum stub for the model. ``run`` reads ``prior_count`` and draws the + # initial particles through ``vector_from_unit_vector``; the prior closure + # isn't traced unless one_step fires. + mock_model = SimpleNamespace( + prior_count=2, + vector_from_unit_vector=lambda v, xp=None: jnp.asarray([0.0, 0.0]), + log_prior_list_from_vector=lambda **kw: [0.0, 0.0], + ) + + resume = MagicMock() + resume.load.return_value = (sentinel_state, [], jax.random.PRNGKey(0), 17) + + paths = NullPaths() + ctx = SimpleNamespace( + model=mock_model, + paths=paths, + resume=resume, + fitness=SimpleNamespace( + objective=lambda kind, compile=None: (lambda params: 0.0), + analysis=SimpleNamespace(), + ), + ) + + search = af.NSS(n_live=4, num_mcmc_steps=1, num_delete=1, termination=-3.0) + + checkpoint = tmp_path / "nss_checkpoint.pkl" + checkpoint.write_bytes(b"resume state") + discarded = [] + search.resume_state = SimpleNamespace( + path=lambda paths: checkpoint, + discard=lambda paths: discarded.append(paths), + ) + + with patch.object( + nss_search_module, + "_blackjax", + ) as mock_bjax, patch.object( + nss_search_module, + "_nss_finalise", + return_value=SimpleNamespace( + particles=SimpleNamespace( + position=np.zeros((1, 2)), + loglikelihood=np.zeros(1), + ), + update_info=SimpleNamespace( + num_expansions=np.zeros(1, dtype=int), + num_shrink=np.zeros(1, dtype=int), + ), + ), + ), patch.object( + nss_search_module, + "_nss_log_weights", + return_value=jnp.zeros((1, 100)), + ): + mock_bjax.nss.return_value.init.side_effect = AssertionError( + "algo.init called from resume path — expected ctx.resume.load instead." + ) + mock_bjax.nss.return_value.step.side_effect = AssertionError( + "algo.step called even though logZ termination should fire immediately." + ) + + internal = search.run(ctx) + + assert isinstance(internal, nss_search_module._NSSInternal) + assert discarded == [] + assert checkpoint.exists() diff --git a/test_autofit/non_linear/test_checkpoint.py b/test_autofit/non_linear/test_checkpoint.py index fe3d207ac..fef4f5d13 100644 --- a/test_autofit/non_linear/test_checkpoint.py +++ b/test_autofit/non_linear/test_checkpoint.py @@ -307,3 +307,54 @@ def test_a_retained_archive_survives_with_search_internal_off( assert "files/search_internal/search_internal.hdf" in _archived(emcee) assert not any("search_internal/" in name for name in _archived(drawer)) + + +class FlakyArchive(DillCheckpointer): + """ + An archive whose first ``finalize`` fails as a full disk would (review 05 finding 2). + """ + + failures = 0 + calls = 0 + + def finalize(self, paths, obj=None): + type(self).calls += 1 + if type(self).failures > 0: + type(self).failures -= 1 + raise OSError("simulated disk full") + super().finalize(paths, obj) + + +class FlakyDrawer(af.Drawer): + checkpointer = FlakyArchive(retain_after_completion=True) + + +def test_a_failed_archive_write_leaves_the_fit_incomplete_and_a_rerun_finishes_it( + paths, monkeypatch +): + """ + ``.completed`` is written only once the archive is finalized: an archive failure + leaves no marker, and the rerun re-finalizes the archive and completes. + """ + import numpy as np + + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + monkeypatch.delenv("PYAUTO_SKIP_FIT_OUTPUT", raising=False) + FlakyArchive.failures = 1 + FlakyArchive.calls = 0 + + analysis = af.ex.Analysis(data=np.full(10, 5.0), noise_map=np.ones(10)) + + first = FlakyDrawer(name="flaky_archive", unique_tag="checkpoint_flaky", total_draws=5) + with pytest.raises(OSError, match="simulated disk full"): + first.fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + assert not first.paths.is_complete + + second = FlakyDrawer(name="flaky_archive", unique_tag="checkpoint_flaky", total_draws=5) + second.fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + assert second.paths.is_complete or Path(second.paths._zip_path).exists() + assert FlakyArchive.calls == 2 + assert "files/search_internal/search_internal.dill" in _archived(second) + assert ".completed" in _archived(second) From 86299f6e8d74036ec06c2604a52edc2750cf70da Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 14 Oct 2026 09:47:32 +0100 Subject: [PATCH 4/6] fix: direct _fit entry builds its PoolFactory from the analysis Outside fit() there is no fit-scoped PoolFactory, and _pools() fell back to PoolFactory(self) without the analysis, so Nautilus(number_of_cores=2) ._fit(model, jax_analysis) saw is_jax=False and took the forking multiprocessing path. _pools(analysis) now applies the JAX fork rule for direct callers; the run(ctx) bridge, Dynesty, the initializer and the sneaky pool pass the analysis they hold. Co-Authored-By: Claude Fable 5.1 --- autofit/non_linear/search/abstract_search.py | 13 +++--- .../search/nest/dynesty/search/abstract.py | 2 +- test_autofit/non_linear/test_fork_rule.py | 46 +++++++++++++++++++ 3 files changed, 54 insertions(+), 7 deletions(-) diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 967291fc4..7823ec64f 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -1521,7 +1521,7 @@ def _fit(self, model: AbstractPriorModel, analysis: Analysis): model=model, analysis=analysis, fitness=fitness, - pool=self._pools(), + pool=self._pools(analysis), test_mode_level=test_mode_level(), ) @@ -1948,7 +1948,7 @@ def start_points( model=model, fitness=fitness, paths=self.paths, - n_cores=self._pools().number_of_cores, + n_cores=self._pools(fitness.analysis).number_of_cores, **initializer_kwargs, ) ) @@ -2054,13 +2054,14 @@ def load_search_internal(self): """ return self.paths.load_search_internal() - def _pools(self) -> PoolFactory: + def _pools(self, analysis: Optional[Analysis] = None) -> PoolFactory: """ The current fit's `PoolFactory` (the JAX fork rule applied to its analysis), or, - outside a fit, one built from `number_of_cores` alone. + outside a fit, one built for `analysis` (so a direct `_fit` caller still gets + the fork rule), or from `number_of_cores` alone when no analysis is given. """ pools = self.__dict__.get("_fit_pools") - return pools if pools is not None else PoolFactory(self) + return pools if pools is not None else PoolFactory(self, analysis) @check_cores def make_pool(self): @@ -2092,7 +2093,7 @@ def make_sneaky_pool(self, fitness: Fitness) -> Optional[SneakyPool]: An implementation of a multiprocessing pool """ - return self._pools().sneaky(fitness) + return self._pools(getattr(fitness, "analysis", None)).sneaky(fitness) def __eq__(self, other): return isinstance(other, NonLinearSearch) and self.__dict__ == other.__dict__ diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index 9a42f3d31..44965920e 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -298,7 +298,7 @@ def _fit( # The fit's effective core count: the JAX fork rule makes it 1 for a JAX # analysis (`autofit.non_linear.parallel.pool`). - number_of_cores = self._pools().number_of_cores + number_of_cores = self._pools(analysis).number_of_cores finished = False diff --git a/test_autofit/non_linear/test_fork_rule.py b/test_autofit/non_linear/test_fork_rule.py index 5ba302a18..9bba2bb9a 100644 --- a/test_autofit/non_linear/test_fork_rule.py +++ b/test_autofit/non_linear/test_fork_rule.py @@ -110,3 +110,49 @@ def test_analysis_pool_refuses_a_jax_analysis(): with pytest.raises(exc.SearchException, match="JAX analyses never fork"): AnalysisPool([_Analysis(is_jax=True)], n_cores=2) + + +class _BackendReached(Exception): + pass + + +def test_direct_fit_entry_applies_the_fork_rule_to_a_jax_analysis(monkeypatch): + """ + Review 05 finding 4: calling the `run(ctx)` bridge `_fit` directly (outside `fit`, + so no fit-scoped `PoolFactory` exists) must still build the context's factory from + the analysis, so Nautilus with a JAX analysis and two cores takes its single-process + path and nothing forks. + """ + pytest.importorskip("jax") + + def no_fork(): + raise AssertionError("must not fork") + + monkeypatch.setattr(pool_module, "fork_context", no_fork) + + search = af.Nautilus(number_of_cores=2) + chosen = {} + + def x1_cpu(fitness, model, update): + chosen["path"] = "fit_x1_cpu" + raise _BackendReached + + def multiprocessing(fitness, model, pools, update): + chosen["path"] = "fit_multiprocessing" + chosen["is_jax"] = pools.is_jax + chosen["cores"] = pools.number_of_cores + raise _BackendReached + + monkeypatch.setattr(search, "fit_x1_cpu", x1_cpu) + monkeypatch.setattr(search, "fit_multiprocessing", multiprocessing) + + analysis = af.ex.Analysis( + data=np.ones(10), noise_map=np.ones(10), use_jax=True + ) + + with pytest.raises(_BackendReached): + search._fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + assert chosen == {"path": "fit_x1_cpu"} + assert search._pools(analysis).is_jax + assert search._pools(analysis).number_of_cores == 1 From 8f35343de17eeca02eee2e7c9464a974dabed10b Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 14 Oct 2026 09:50:33 +0100 Subject: [PATCH 5/6] fix: FitContext closes its pools on success and ctx.update saves the checkpoint The frozen run(ctx) contract promises cleanup on every exit and a checkpoint save on each update. ctx.close(failed=False) returned at once, leaking context-built pools after a successful run, and ctx.update only called perform_update, so a resumable run(ctx) search lost its progress on interruption. PoolFactory.close(failed) now closes+joins on success and terminates+joins on failure (idempotent); ctx.update saves the search's resume_state (never finalize) before the update. Co-Authored-By: Claude Fable 5.1 --- autofit/non_linear/parallel/pool.py | 27 ++-- autofit/non_linear/search/fit_context.py | 26 ++- docs/design/run_ctx.md | 10 +- .../non_linear/search/test_fit_context.py | 149 ++++++++++++++++++ 4 files changed, 195 insertions(+), 17 deletions(-) diff --git a/autofit/non_linear/parallel/pool.py b/autofit/non_linear/parallel/pool.py index c9ad6cd65..05d5fb101 100644 --- a/autofit/non_linear/parallel/pool.py +++ b/autofit/non_linear/parallel/pool.py @@ -222,17 +222,26 @@ def sneaky(self, fitness): self._created.append(pool) return pool - def close(self): + def close(self, failed: bool = True): """ - Terminate every pool this factory built, after a failed fit. A successful fit's - pools are closed by the search that used them, exactly as before. + Shut down every pool this factory built, once: ``close`` then ``join`` after a + successful fit (the workers finish and exit), ``terminate`` then ``join`` after a + failed one. Idempotent, and safe on a pool its backend already shut down. """ - for pool in self._created: - terminate = getattr(pool, "terminate", None) - if terminate is None: + created, self._created = self._created, [] + for pool in created: + for name in ("terminate",) if failed else ("close",): + method = getattr(pool, name, None) + if method is None: + continue + try: + method() + except Exception: + logger.debug("Pool %s after a fit raised", name, exc_info=True) + join = getattr(pool, "join", None) + if join is None: continue try: - terminate() + join() except Exception: - logger.debug("Pool termination after a failed fit raised", exc_info=True) - self._created = [] + logger.debug("Pool join after a fit raised", exc_info=True) diff --git a/autofit/non_linear/search/fit_context.py b/autofit/non_linear/search/fit_context.py index 9f2c29dd6..a557b7ad8 100644 --- a/autofit/non_linear/search/fit_context.py +++ b/autofit/non_linear/search/fit_context.py @@ -121,7 +121,19 @@ def update(self, internal: Any): The backend is paused for the duration of the call (the update is synchronous), so ``internal`` is read as it stands at this point of the run. + + The checkpoint is saved first: the search's resume state (``ctx.resume``'s + strategy, ``NonLinearSearch.resume_state``) gets ``save(internal)``, so an + interrupted run resumes from the last update. A native resume state (one the + backend writes itself) keeps its no-op ``save``; a search without a resume + state is not resumable and saves nothing until its archive is finalized at + completion. ``finalize`` is never called here, since that is the completion step + that may discard the resume state. """ + resume_state = self._search.resume_state + if resume_state is not None: + resume_state.save(self.paths, internal) + return self._search.perform_update( model=self.model, analysis=self._analysis, @@ -132,15 +144,19 @@ def update(self, internal: Any): def close(self, failed: bool): """ - Clean up after ``run``. On failure: close the pools the context built, shut down - the quick-update threads and release the compiled objectives. On success the - fitness is still needed for the final update, which shuts its threads down. + Clean up after ``run``, on every exit. The pools the context built are shut + down on both paths (closed and joined on success, terminated and joined on + failure; ``PoolFactory.close``). On failure the quick-update threads are also + shut down and the compiled objectives released; on success the fitness is still + needed for the final update, which shuts its threads down. """ + close_pools = getattr(self.pool, "close", None) + if close_pools is not None: + close_pools(failed=failed) + if not failed: return - self.pool.close() - shutdown = getattr(self.fitness, "shutdown_quick_update", None) if shutdown is not None: shutdown() diff --git a/docs/design/run_ctx.md b/docs/design/run_ctx.md index 2b30ff440..5c57a870d 100644 --- a/docs/design/run_ctx.md +++ b/docs/design/run_ctx.md @@ -136,9 +136,13 @@ pins the following details the table left open; each is flagged for review in th `Checkpointer` handle, not the loaded state, because a backend may reopen its own native file (`Nautilus` resumes through `Sampler(filepath=...)`). **`ctx.rng`** is a `SeedSequence` of the search's `seed` attribute, if any, until A4's search-level seed. -- **`ctx.close(failed)`** is the cleanup the lifecycle rules require; on failure it closes - the pools the context built, shuts the quick-update threads down and releases the - compiled objectives. +- **`ctx.close(failed)`** is the cleanup the lifecycle rules require. On every exit it + shuts down the pools the context built (closed and joined on success, terminated and + joined on failure); on failure it also shuts the quick-update threads down and + releases the compiled objectives. +- **`ctx.update(internal)`** saves the checkpoint before the update: the search's + `resume_state` gets `save(internal)` (a no-op for a native file the backend writes + itself), never `finalize`, which is the completion step. - **The bridge is `NonLinearSearch._fit`**: when a search defines `run`, the base `_fit` builds the `Fitness` and the context (after the gates `start_resume_fit` runs) and calls `run(ctx)`, so every caller of `_fit` (and the conformance suite's call counter) sees a diff --git a/test_autofit/non_linear/search/test_fit_context.py b/test_autofit/non_linear/search/test_fit_context.py index 85ad6159a..fdc18ad0f 100644 --- a/test_autofit/non_linear/search/test_fit_context.py +++ b/test_autofit/non_linear/search/test_fit_context.py @@ -187,3 +187,152 @@ def ctx(): path.write_bytes(b"state") assert ctx().resume is search.resume_state + + +# Review 05 findings 5 and 6: the frozen contract promises that the context closes the +# pools it built on every exit, and that `ctx.update` saves the checkpoint. + + +class _RecordingPool: + def __init__(self, calls, processes, **kwargs): + self.calls = calls + calls.append("created") + + def close(self): + self.calls.append("close") + + def terminate(self): + self.calls.append("terminate") + + def join(self): + self.calls.append("join") + + +@pytest.fixture(name="pool_calls") +def make_pool_calls(monkeypatch): + from types import SimpleNamespace + + from autofit.non_linear.parallel import pool as pool_module + + calls = [] + monkeypatch.setattr( + pool_module, + "fork_context", + lambda: SimpleNamespace( + Pool=lambda processes, **kwargs: _RecordingPool(calls, processes, **kwargs) + ), + ) + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + return calls + + +class PoolToySearch(ToySearch): + fail = False + + def run(self, ctx): + assert ctx.pool() is not None + internal = super().run(ctx) + if self.fail: + raise RuntimeError("backend failed") + return internal + + +def test_a_successful_run_closes_and_joins_the_pools_the_context_built(pool_calls): + search = PoolToySearch(number_of_cores=2, initializer=af.InitializerPrior()) + search.fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + assert pool_calls == ["created", "close", "join"] + + +def test_a_failed_run_terminates_and_joins_the_pools_the_context_built(pool_calls): + search = PoolToySearch(number_of_cores=2, initializer=af.InitializerPrior()) + search.fail = True + + with pytest.raises(RuntimeError, match="backend failed"): + search.fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + assert pool_calls == ["created", "terminate", "join"] + + +def test_pool_factory_close_is_idempotent(pool_calls): + from autofit.non_linear.parallel import PoolFactory + + factory = PoolFactory(ToySearch(number_of_cores=2)) + factory() + factory.close(failed=False) + factory.close(failed=True) + + assert pool_calls == ["created", "close", "join"] + + +def _recording_resume_state(): + from autofit.non_linear.checkpoint import DillCheckpointer + + class RecordingCheckpointer(DillCheckpointer): + def __init__(self): + super().__init__(filename="toy_resume.dill") + self.saved = [] + self.finalized = [] + + def save(self, paths, obj): + self.saved.append(obj) + super().save(paths, obj) + + def finalize(self, paths, obj=None): + self.finalized.append(obj) + + return RecordingCheckpointer() + + +class UpdatingToySearch(ToySearch): + interrupt = False + + def run(self, ctx): + internal = super().run(ctx) + ctx.update(internal) + if self.interrupt: + raise InterruptedError("killed after the update") + return internal + + +def test_ctx_update_saves_the_resume_state_so_an_interrupted_run_keeps_it( + tmp_path, monkeypatch +): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + strategy = _recording_resume_state() + + class Resumable(UpdatingToySearch): + resume_state = strategy + + search = Resumable( + initializer=af.InitializerPrior(), path_prefix=str(tmp_path), name="toy_update" + ) + search.interrupt = True + + with pytest.raises(InterruptedError): + search.fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + assert len(strategy.saved) == 1 + assert len(strategy.saved[0]["log_likelihoods"]) == 6 + # Saved through `save`, never the completion `finalize`, and still on disk to resume. + assert strategy.finalized == [] + assert strategy.exists(search.paths) + + +def test_ctx_update_saves_nothing_for_a_search_without_a_resume_state( + tmp_path, monkeypatch +): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + saved = [] + search = UpdatingToySearch( + initializer=af.InitializerPrior(), path_prefix=str(tmp_path), name="toy_plain" + ) + monkeypatch.setattr( + type(search.checkpointer), "save", lambda self, paths, obj: saved.append(obj) + ) + search.fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + # The archive is written once, at completion, not by the mid-run update. + assert len(saved) <= 1 From 2090ca9dd0595fdd5d87b7d1ebce8f2fdd46731e Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 14 Oct 2026 09:52:10 +0100 Subject: [PATCH 6/6] test: pin the golden BlackJAX diagnostics by value, reference and ownership test_samples_info_is_unchanged compared ESS / R-hat by shape only, so scalar diagnostics had no protection and a shape-preserving -999 passed. They are now compared by value to the fixture (rtol 0.25, finite, positive), exactly against the installed blackjax estimators on the stored chains, and via deterministic estimator doubles that must be forwarded unchanged; a parametrized sentinel test proves every key fails on -999. The samples.csv tolerance (rtol 1e-12) is unchanged. Co-Authored-By: Claude Fable 5.1 --- .../non_linear/samples/test_golden_samples.py | 136 +++++++++++++++++- 1 file changed, 132 insertions(+), 4 deletions(-) diff --git a/test_autofit/non_linear/samples/test_golden_samples.py b/test_autofit/non_linear/samples/test_golden_samples.py index eb18f270f..4430e0e16 100644 --- a/test_autofit/non_linear/samples/test_golden_samples.py +++ b/test_autofit/non_linear/samples/test_golden_samples.py @@ -19,7 +19,15 @@ ``samples_info`` is compared exactly except for the convergence diagnostics BlackJAX itself computes (ESS, bulk/tail ESS, rank-normalised R-hat), whose estimators change -between BlackJAX releases; for those only the shape is pinned. +between BlackJAX releases. Those are pinned three ways, so a wrong value or wrong +wiring cannot pass on shape alone (review 05 finding 7): + +- by value against the fixture, to ``DIAGNOSTIC_RTOL`` (loose enough for an estimator + revision, far tighter than any wiring error), finite and positive; +- exactly against the installed BlackJAX's own calculation on the stored chains + (``test_diagnostics_are_the_installed_blackjax_reference``); +- by ownership, with the BlackJAX estimators replaced by deterministic doubles whose + outputs must be forwarded unchanged (``test_diagnostics_forward_the_backend_outputs``). """ import csv import importlib.util @@ -67,6 +75,32 @@ def reference_mode(monkeypatch): ) +#: Relative tolerance on the BlackJAX diagnostics against the fixture: room for an +#: estimator revision between BlackJAX releases, while a sentinel, a swapped key or a +#: per-parameter mix-up moves them by far more. +DIAGNOSTIC_RTOL = 0.25 + + +def assert_diagnostics_match(info: dict, expected: dict): + """ + Pop every ``BACKEND_DIAGNOSTIC_KEYS`` entry from ``info`` and ``expected`` and + compare them by value: the same shape, finite, positive (ESS and R-hat both are) and + within ``DIAGNOSTIC_RTOL`` of the fixture. + """ + for key in BACKEND_DIAGNOSTIC_KEYS: + if key not in expected: + continue + actual = np.asarray(info.pop(key), dtype=float) + reference = np.asarray(expected.pop(key), dtype=float) + + assert actual.shape == reference.shape, key + assert np.all(np.isfinite(actual)), key + assert np.all(actual > 0.0), key + np.testing.assert_allclose( + actual, reference, rtol=DIAGNOSTIC_RTOL, atol=0.0, err_msg=key + ) + + def _read_csv(path): with open(path, newline="") as f: rows = list(csv.reader(f)) @@ -101,8 +135,102 @@ def test_samples_info_is_unchanged(name): info = cases.samples_info_for_comparison(samples) - for key in BACKEND_DIAGNOSTIC_KEYS: - if key in expected: - assert np.shape(info.pop(key)) == np.shape(expected.pop(key)), key + assert_diagnostics_match(info, expected) assert info == expected + + +# Review 05 finding 7: the diagnostics must be protected by value, not only by shape. + + +def _sentinel(value): + return (np.full(np.shape(value), -999.0)).tolist() if np.ndim(value) else -999.0 + + +@pytest.mark.parametrize("key", BACKEND_DIAGNOSTIC_KEYS) +def test_sentinel_diagnostics_fail_the_golden_comparison(key, monkeypatch): + """ + Replacing any backend diagnostic with a shape-preserving ``-999`` must fail + ``test_samples_info_is_unchanged``. + """ + _skip_missing("nuts") + + original = cases.samples_info_for_comparison + + def mutated(samples): + info = original(samples) + info[key] = _sentinel(info[key]) + return info + + monkeypatch.setattr(cases, "samples_info_for_comparison", mutated) + + with pytest.raises(AssertionError): + test_samples_info_is_unchanged("nuts") + + +def test_diagnostics_are_the_installed_blackjax_reference(): + """ + The NUTS diagnostics equal the installed BlackJAX's own estimators run on the stored + chains, reshaped chain-major exactly as the search documents. + """ + _skip_missing("nuts") + import jax.numpy as jnp + from blackjax import diagnostics + + from autofit.non_linear.search.mcmc.blackjax.nuts.search import ( + split_chain_diagnostics, + ) + + internal = cases.internal_from("nuts") + info = cases.samples_info_for_comparison(cases.samples_from("nuts")) + + chain_major = jnp.asarray(split_chain_diagnostics(internal["positions"])) + + reference = { + "ess_per_param": diagnostics.effective_sample_size(chain_major), + "ess_bulk_per_param": diagnostics.ess_bulk(chain_major), + "ess_tail_per_param": diagnostics.ess_tail(chain_major), + "rhat_per_param": diagnostics.rhat(chain_major), + } + + for key, value in reference.items(): + value = np.atleast_1d(np.asarray(value)) + np.testing.assert_allclose(info[key], value, rtol=1e-12, atol=0.0, err_msg=key) + + assert info["ess_min"] == pytest.approx(min(info["ess_per_param"]), rel=1e-15) + assert info["ess_bulk_min"] == pytest.approx(min(info["ess_bulk_per_param"]), rel=1e-15) + assert info["ess_tail_min"] == pytest.approx(min(info["ess_tail_per_param"]), rel=1e-15) + assert info["rhat_max"] == pytest.approx(max(info["rhat_per_param"]), rel=1e-15) + + +def test_diagnostics_forward_the_backend_outputs(monkeypatch): + """ + Ownership: with each BlackJAX estimator replaced by a deterministic double, the + search forwards its output unchanged to the matching ``samples_info`` key and takes + the documented min / max. + """ + _skip_missing("nuts") + from blackjax import diagnostics + + doubles = { + "effective_sample_size": [11.0, 7.0, 13.0], + "ess_bulk": [21.0, 17.0, 23.0], + "ess_tail": [31.0, 37.0, 33.0], + "rhat": [1.01, 1.07, 1.03], + } + for name, value in doubles.items(): + monkeypatch.setattr( + diagnostics, name, lambda chains, _value=value: np.asarray(_value) + ) + + internal = cases.internal_from("nuts") + info = cases.search_from("nuts").samples_info_from(internal) + + assert info["ess_per_param"] == doubles["effective_sample_size"] + assert info["ess_min"] == 7.0 + assert info["ess_bulk_per_param"] == doubles["ess_bulk"] + assert info["ess_bulk_min"] == 17.0 + assert info["ess_tail_per_param"] == doubles["ess_tail"] + assert info["ess_tail_min"] == 31.0 + assert info["rhat_per_param"] == doubles["rhat"] + assert info["rhat_max"] == 1.07