diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 1bcebb7dc..b5528b393 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -402,6 +402,9 @@ def __init__( test_mode_budget = {} objective_target = None invalid_value = -float("inf") + # True for a search whose backend is wrong in float32: `fit` then raises, rather + # than warns, when a JAX analysis runs with `jax_enable_x64` off (A3b preflight). + requires_fp64 = False def optimise( self, @@ -899,6 +902,16 @@ def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Res # smoke runs use `PYAUTO_TEST_MODE>=2`, which returned above. cap.check_jax_required(self, analysis) + # The JAX checks (search-extensibility A3b, `autofit.non_linear.search.preflight`), + # after the test-mode bypass and the REQUIRED gate and before any backend state + # exists: warn once when x64 is off (raise for `requires_fp64`), then trace the + # declared objective kind(s) once with `jax.eval_shape` so a non-traceable + # likelihood fails here, naming the search, rather than inside the backend. + from autofit.non_linear.search import preflight + + preflight.check_x64(self, analysis) + preflight.trace_preflight(self, analysis, model) + # The fit's pools, with the one JAX fork rule applied once (one INFO line when it # downgrades) and the effective worker count recorded for `search.summary`. self._fit_pools = PoolFactory(self, analysis) @@ -1458,6 +1471,10 @@ def _fit(self, model: AbstractPriorModel, analysis: Analysis): from autofit.non_linear.search.fit_context import FitContext + # The REQUIRED gate, repeated for callers that reach `_fit` without `fit`, before + # any `Fitness` or backend state exists (`fit` has already run it). + cap.check_jax_required(self, analysis) + fitness = self.make_fitness( analysis=analysis, model=model, **self.fitness_overrides(analysis) ) diff --git a/autofit/non_linear/search/capabilities.py b/autofit/non_linear/search/capabilities.py index 283835461..85b73e781 100644 --- a/autofit/non_linear/search/capabilities.py +++ b/autofit/non_linear/search/capabilities.py @@ -32,6 +32,9 @@ ``checkpointer`` Where the search archives its internal state when a fit completes (``autofit.non_linear.checkpoint``): ``dill``, ``pickle`` or ``native:``. +``requires_fp64`` Whether the backend needs float64: ``fit`` then raises, rather + than warns once, when a JAX analysis runs with + ``jax_enable_x64`` off (``autofit.non_linear.search.preflight``). ======================== ============================================================== ``invalid_value`` records what the backend actually sees and is not yet normalised. For @@ -176,6 +179,7 @@ def to_dict(self) -> Dict[str, str]: "objective_target", "invalid_value", "checkpointer", + "requires_fp64", ) """ The names of the static capability attributes, in manifest order. diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index e6328ec51..f69232412 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -4,6 +4,7 @@ import os import sys from abc import ABC +from contextlib import ExitStack from pathlib import Path from typing import Dict, Optional, Tuple, Union, TYPE_CHECKING @@ -55,6 +56,21 @@ def __enter__(self): return ForkPool +class _SingleCoreRun(Exception): + """ + Control flow inside ``AbstractDynesty._fit``: run dynesty without a pool. + + Raised when the fit is single-core by construction (``number_of_cores <= 1``, + ``force_x1_cpu`` or a JAX analysis) and when the multiprocessing pool cannot be + created (``RuntimeError`` from its ``__enter__``, e.g. on platforms with weak + multiprocessing support), which is then its ``__cause__``. It is a private + ``Exception`` subclass, not a ``RuntimeError``, so the ``except`` that catches it + never swallows a ``RuntimeError`` raised while dynesty runs (a genuine dynesty + error, or JAX's ``XlaRuntimeError``, which subclasses ``RuntimeError``): those + propagate instead of silently restarting the run single-core. + """ + + DYNESTY_FILENAME = "savestate.save" """ dynesty's checkpoint file, the search's resume state. @@ -234,8 +250,9 @@ def _fit( as on the existing `force_x1_cpu` path. However, certain operating systems (e.g. Windows) do not support Python multiprocessing particularly well. - This can cause Dynesty to crash when a pool is included. If this occurs (raising a `RunTimeException`) - a Dynesty object without a pool is created and used instead. + If creating the pool raises a `RuntimeError`, a Dynesty object without a pool is created and used instead. + Only the pool's creation falls back: a `RuntimeError` raised while Dynesty runs (including JAX's + `XlaRuntimeError`) propagates rather than silently restarting the run single-core (`_SingleCoreRun`). Parameters ---------- @@ -283,17 +300,26 @@ def _fit( while not finished: try: if number_of_cores <= 1 or self.force_x1_cpu or analysis.is_jax: - raise RuntimeError + raise _SingleCoreRun() Pool = _fork_pool_cls() - with Pool( - njobs=number_of_cores, - loglike=fitness, - prior_transform=prior_transform, - logl_args=(model, fitness), - ptform_args=(model,), - ) as pool: + with ExitStack() as stack: + # Only a failure to *create* the pool falls back to a single-core + # run; a RuntimeError raised while dynesty runs propagates. + try: + pool = stack.enter_context( + Pool( + njobs=number_of_cores, + loglike=fitness, + prior_transform=prior_transform, + logl_args=(model, fitness), + ptform_args=(model,), + ) + ) + except RuntimeError as error: + raise _SingleCoreRun() from error + search_internal = self.search_internal_from( model=model, fitness=fitness, @@ -306,7 +332,7 @@ def _fit( checkpoint_exists = True - except RuntimeError as e: + except _SingleCoreRun as e: if not checkpoint_exists: if analysis.is_jax: self.logger.info( @@ -323,7 +349,7 @@ def _fit( else: self.logger.info( f""" - The Dynesty multiprocessing pool could not be created ({e!r}). + The Dynesty multiprocessing pool could not be created ({e.__cause__!r}). A single CPU non-multiprocessing Dynesty run is being performed. """ diff --git a/autofit/non_linear/search/nest/nss/_chunked_nss.py b/autofit/non_linear/search/nest/nss/_chunked_nss.py index b7f360392..2586977a8 100644 --- a/autofit/non_linear/search/nest/nss/_chunked_nss.py +++ b/autofit/non_linear/search/nest/nss/_chunked_nss.py @@ -25,7 +25,7 @@ When ``chunk_size`` is None the builder still uses ``jax.vmap`` and is bit-identical to upstream. The builder otherwise produces a ``blackjax.SamplingAlgorithm`` with the same shape ``blackjax.nss(...)`` -returns, so ``af.NSS._fit`` is a one-line switch. +returns, so ``af.NSS.run`` is a one-line switch. See PyAutoFit#1304 for the diagnosis and A100 evidence (jobs 322605 / 322606 OOM at the same byte counts as before #1303 landed, because the diff --git a/autofit/non_linear/search/nest/nss/search.py b/autofit/non_linear/search/nest/nss/search.py index 2856e34fb..e2b5f5c3b 100644 --- a/autofit/non_linear/search/nest/nss/search.py +++ b/autofit/non_linear/search/nest/nss/search.py @@ -8,10 +8,8 @@ import numpy as np -from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.checkpoint import NativeFileCheckpointer -from autofit.non_linear.fitness import Fitness, get_log_likelihood_ceiling -from autofit.non_linear.paths.null import NullPaths +from autofit.non_linear.objective import SCALAR from autofit.non_linear.search.nest import abstract_nest from .samples import NSSamples from autofit.non_linear.samples.adapter import RawSamples @@ -41,70 +39,15 @@ logger = logging.getLogger(__name__) -#: The figure of merit ``af.NSS`` substitutes for a log likelihood it refuses to sample from. Not -#: ``-inf``: blackjax's nested sampler arithmetic (log weights, shell evidence) has to stay finite, -#: so an unreachably low finite value is used instead. +#: The figure of merit ``af.NSS`` substitutes for a log likelihood it refuses to sample from +#: (its declared ``invalid_value``). Not ``-inf``: blackjax's nested sampler arithmetic (log +#: weights, shell evidence) has to stay finite, so an unreachably low finite value is used +#: instead. ``Fitness.call`` returns it for a non-finite log likelihood, for one over the +#: configured magnitude ceiling and, through the traced ``xp.where`` penalty, for a model +#: that violates an assertion. NSS_INVALID_LOG_LIKELIHOOD = -1e30 -def nss_log_likelihood_from(model, analysis, log_likelihood_ceiling=None): - """ - Build the JAX log-likelihood closure ``af.NSS`` samples through. - - ``af.NSS`` does not call `Fitness.call` while sampling -- it hands blackjax an inline JAX - closure -- so it does not inherit that method's guards and has to repeat them here. Keeping the - closure in one module-level factory means the guards exist once and can be tested directly, - rather than being restated by a test that only resembles the sampled path. - - Two things are rejected, both mapped to `NSS_INVALID_LOG_LIKELIHOOD`: - - - non-finite values (``NaN`` / ``inf``), always; - - finite values whose magnitude exceeds the configured ceiling, **if** a config has opted in. - An fp64 Cholesky on a non-positive-definite matrix returns finite garbage (log likelihoods up - to ``3e+303``), which a nested sampler otherwise accepts as its highest-likelihood live point; - the shell log evidence then explodes and the termination criterion never fires. The ceiling is - off in the packaged config, because a log likelihood scales with the noise-map units and a - fixed magnitude can therefore reject a legitimate fit - (`autofit.non_linear.fitness.get_log_likelihood_ceiling` carries the full argument). - - This closure is JAX-only, so it cannot warn when the ceiling fires — the rejection is an - ``jnp.where`` on a tracer. `Fitness.call`'s numpy path warns once per process; ``af.NSS`` does - not sample through it. - - Parameters - ---------- - model - The model whose parameter vector is mapped to an instance. - analysis - The analysis supplying `log_likelihood_function`. - log_likelihood_ceiling - The magnitude ceiling. Defaults to the configured value - (`autofit.non_linear.fitness.get_log_likelihood_ceiling`), which is ``inf`` — the check - disabled — unless a config opts in. It must be a static Python float, because JAX traces - the comparison it feeds. - - Returns - ------- - The closure blackjax samples through. - """ - import jax.numpy as jnp - - if log_likelihood_ceiling is None: - log_likelihood_ceiling = get_log_likelihood_ceiling() - - def log_likelihood(params): - instance = model.instance_from_vector(vector=params, xp=jnp) - raw = analysis.log_likelihood_function(instance=instance) - raw = jnp.where(jnp.isfinite(raw), raw, NSS_INVALID_LOG_LIKELIHOOD) - return jnp.where( - jnp.abs(raw) > log_likelihood_ceiling, - NSS_INVALID_LOG_LIKELIHOOD, - raw, - ) - - return log_likelihood - - _CHECKPOINT_FILENAME = "nss_checkpoint.pkl" @@ -216,7 +159,7 @@ class NSS(abstract_nest.AbstractNest): # physical-space prior density. objective_target = cap.ObjectiveTarget(cap.ObjectiveQuantity.LOG_LIKELIHOOD, cap.CoordinateSpace.PHYSICAL) invalid_value = NSS_INVALID_LOG_LIKELIHOOD - # The resume state is NSS's own atomic ``nss_checkpoint.pkl`` (``_fit`` writes, + # The resume state is NSS's own atomic ``nss_checkpoint.pkl`` (``run`` writes, # reads and deletes it); the archive is the default ``search_internal.dill``. resume_state = NativeFileCheckpointer(_CHECKPOINT_FILENAME, loader=_load_checkpoint) @@ -397,47 +340,59 @@ def apply_test_mode(self): ) self.termination = -1.0 - def _fit(self, model: AbstractPriorModel, analysis): + def fitness_overrides(self, analysis) -> dict: """ - Fit a model using NSS, with checkpoint/resume + on-the-fly visualization. - - Builds JAX-traceable ``log_likelihood`` and ``prior_logprob`` closures - threaded through Phase 0's ``xp=jnp`` plumbing, draws ``n_live`` initial - particles by mapping unit-cube samples through the prior transform, - and runs the NSS outer loop inline (mirroring upstream blackjax's - nested-sampling loop pattern). Between outer iterations the - loop can (a) pickle resumable state to ``nss_checkpoint.pkl`` and - (b) call ``analysis.visualize`` on the current best live particle. - - On entry, if a checkpoint exists at the expected path the loop resumes - from the saved ``(state, dead, run_key, iteration)``. On successful - exit the checkpoint is deleted — mirrors Nautilus's - ``output_search_internal`` post-success cleanup so the next fresh fit - doesn't accidentally resume from a stale checkpoint. - - Returns - ------- - (search_internal, fitness) - ``search_internal`` is a ``_NSSInternal`` holder (NumPy arrays - only). ``fitness`` is a ``Fitness`` instance that ``af.NSS`` does - not use for sampling (inline JAX closures handle that) but is - required by ``AbstractNest.perform_update`` for post-fit work - like latent-sample generation, which calls ``fitness.batch_size``. + NSS fires its own quick update between outer iterations (``_fire_quick_update``), + so the ``Fitness`` cadence that ``call_wrap`` drives is switched off: blackjax + evaluates the raw objective inside its jitted step and never goes through + ``call_wrap``, and a second background worker would only duplicate the visuals. """ + return { + "iterations_per_quick_update": None, + "background_quick_update": False, + "live_visual_update": False, + } - # jax_use='required': NSS traces its likelihood inside blackjax and has no numpy - # path. `fit` already ran this gate; repeated for direct callers. - cap.check_jax_required(self, analysis) + def run(self, ctx): + """ + Run NSS (the ``run(ctx)`` hook; see ``docs/design/run_ctx.md``) and return the + ``_NSSInternal`` holder of its final state. + + The log likelihood blackjax samples through is the fit's scalar objective + (``Fitness.call`` via ``Fitness.objective``), so NSS shares the one guard path + every search uses: a non-finite log likelihood, one over the configured + magnitude ceiling, and a model that violates an assertion (the traced + ``xp.where`` penalty) all return the declared ``invalid_value`` + (``NSS_INVALID_LOG_LIKELIHOOD``). It is handed over unjitted: blackjax traces it + inside its own jitted step, where a nested jit would only add a call boundary. + + The prior log density and the unit-cube initial draws are built from + ``ctx.model`` with ``xp=jax.numpy``. Between outer iterations the loop (a) + pickles resumable state to NSS's native ``nss_checkpoint.pkl`` (the search's + ``resume_state``) and (b) fires a quick-update visualization of the current best + live particle. When ``ctx.resume`` holds that file the loop resumes from the + saved ``(state, dead, run_key, iteration)``; on successful exit the file is + deleted, so the next fresh fit does not resume from a stale checkpoint. + + The likelihood-function sanity check of a resumed run runs when the fit's + ``Fitness`` is built, before this method, rather than after sampling. + Parameters + ---------- + ctx + The fit's ``FitContext``. + """ import jax import jax.numpy as jnp import time + model = ctx.model + self.logger.info("Starting NSS non-linear search.") ndim = model.prior_count - log_likelihood = nss_log_likelihood_from(model=model, analysis=analysis) + log_likelihood = ctx.fitness.objective(SCALAR, compile=False) def prior_logprob(params): log_priors = model.log_prior_list_from_vector(vector=params, xp=jnp) @@ -493,9 +448,10 @@ def one_step(carry, _): state, dead_point = algo.step(subk, state) return (state, k), dead_point - checkpoint_path = self._nss_checkpoint_path - if checkpoint_path is not None and checkpoint_path.exists(): - state, dead, run_key, iteration = _load_checkpoint(checkpoint_path) + checkpoint_path = self.resume_state.path(ctx.paths) + + if ctx.resume is not None: + state, dead, run_key, iteration = ctx.resume.load(ctx.paths) self.logger.info( "Resuming NSS from checkpoint at iteration %d (state file %s).", iteration, @@ -534,7 +490,9 @@ def one_step(carry, _): self.iterations_per_quick_update is not None and iteration % self.iterations_per_quick_update == 0 ): - self._fire_quick_update(state=state, model=model, analysis=analysis) + self._fire_quick_update( + state=state, model=model, analysis=ctx.fitness.analysis + ) wall_time = time.time() - t_start @@ -580,7 +538,7 @@ def _safe_ess(log_w_mean): if checkpoint_path is not None and checkpoint_path.exists(): try: - checkpoint_path.unlink() + self.resume_state.discard(ctx.paths) except OSError as exc: self.logger.warning( "Failed to delete completed-run checkpoint %s: %s. The " @@ -590,34 +548,21 @@ def _safe_ess(log_w_mean): exc, ) - fitness = Fitness( - model=model, - analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=True, - resample_figure_of_merit=-1.0e99, - batch_size=1, - ) - - return search_internal, fitness + return search_internal @property def _nss_checkpoint_path(self) -> Optional[Path]: - """Resolve the checkpoint location, or None when paths is NullPaths.""" - if isinstance(self.paths, NullPaths): - return None - try: - return Path(self.paths.search_internal_path) / _CHECKPOINT_FILENAME - except TypeError: - return None + """Resolve the checkpoint location (the search's ``resume_state`` file), or None + when the paths store no internal state (``NullPaths``).""" + return self.resume_state.path(self.paths) def _fire_quick_update(self, state, model, analysis): """Push the current best live particle through ``analysis.visualize``. The Nautilus / Dynesty quick-update path goes through - ``Fitness.manage_quick_update``; ``af.NSS`` bypasses ``Fitness._call`` - for sampling so we invoke ``analysis.visualize`` directly between - outer-loop iterations. Wrapped in try/except — a visualization failure + ``Fitness.manage_quick_update``; ``af.NSS`` hands blackjax the raw + objective (no ``call_wrap`` bookkeeping) so we invoke + ``analysis.visualize`` directly between outer-loop iterations. Wrapped in try/except — a visualization failure logs a warning but does not kill a long sampler run. """ try: diff --git a/autofit/non_linear/search/preflight.py b/autofit/non_linear/search/preflight.py new file mode 100644 index 000000000..3b31e0b4d --- /dev/null +++ b/autofit/non_linear/search/preflight.py @@ -0,0 +1,227 @@ +""" +The JAX checks ``NonLinearSearch.fit`` runs before a search's backend builds any state +(search-extensibility phase A3b; survey 02 §5.3, decisions D2 and D8). + +They run in ``NonLinearSearch.start_resume_fit`` after the test-mode bypass return (so a +``PYAUTO_TEST_MODE>=2`` smoke run, which never runs a backend, is unaffected) and after +the ``jax_use='required'`` gate (``capabilities.check_jax_required``), and only when the +analysis is JAX (``Analysis.is_jax``): + +- ``check_x64``: JAX defaults to 32-bit floats, and x64 is switched on only through an + 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 + 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 + array) fails here, naming the search and chaining the original tracer error, instead + of deep inside the backend after its state exists. It is a **trace check, not a + validity certificate**: a likelihood that traces can still return NaN. +- the optional numerical probe (``PYAUTO_JAX_PREFLIGHT_PROBE=1``, off by default): + evaluate the scalar objective once, eagerly, at the prior medians and warn when it + returns the search's invalid value. It executes the likelihood, so it is never on by + default. + +``PYAUTO_JAX_PREFLIGHT=0`` skips the trace preflight (the escape hatch for a likelihood +deliberately run eagerly, e.g. ``compile=False`` debugging of a function that +does not trace). + +JAX is imported inside the functions, so importing this module loads no optional +backend. +""" + +import logging +import os +import warnings + +from autofit import exc +from autofit.non_linear.objective import ( + BATCHED, + BATCHED_VALUE_AND_GRAD, + SCALAR, + VALUE_AND_GRAD, +) +from autofit.non_linear.search import capabilities as cap + +logger = logging.getLogger(__name__) + +PREFLIGHT_ENV = "PYAUTO_JAX_PREFLIGHT" +""" +Set to ``0`` (or ``false`` / ``no`` / ``off``) to skip the trace preflight. +""" + +PROBE_ENV = "PYAUTO_JAX_PREFLIGHT_PROBE" +""" +Set to ``1`` (or ``true`` / ``yes`` / ``on``) to run the numerical probe. +""" + +X64_WARNING = ( + "JAX is running in 32-bit precision (jax_enable_x64 is False), so the likelihood " + "of this JAX analysis is evaluated in float32. PyAutoFit enables x64 through an " + "environment variable read when autonerves is imported; importing jax before " + "autofit (or setting JAX_ENABLE_X64=0) leaves it off. Call " + "jax.config.update('jax_enable_x64', True) before building the analysis, or import " + "autofit before jax, to fit in float64." +) +""" +The warning ``check_x64`` emits once per process. +""" + +X64_REQUIRED_MESSAGE = ( + "{search} requires float64 (requires_fp64=True), but JAX is running in 32-bit " + "precision (jax_enable_x64 is False). Call jax.config.update('jax_enable_x64', " + "True) before building the analysis, or import autofit before jax." +) +""" +The error ``check_x64`` raises for a search that declares ``requires_fp64``. +""" + +PREFLIGHT_MESSAGE = ( + "{search} cannot run on this JAX analysis: tracing its {kind!r} objective at the " + "prior-median parameters failed with {error_type}: {error}\n\n" + "This is the trace preflight, run before the search builds any state: the " + "likelihood of {analysis} (use_jax=True) must be traceable by jax.jit. Replace numpy " + "calls on parameter-dependent values (np.asarray, np.array, float(), a Python if on " + "an array) with jax.numpy / xp equivalents, or construct the analysis with " + "use_jax=False and choose a search whose jax_use is not 'required'. The original " + "error is chained above. Set {env}=0 to skip this check." +) +""" +The error ``trace_preflight`` raises, chained to the tracer error. +""" + +_x64_warning_emitted = False + + +def _env_flag(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() not in ("0", "false", "no", "off", "") + + +def check_x64(search, analysis): + """ + Warn once per process when a JAX analysis runs with ``jax_enable_x64`` off, or + raise ``SearchException`` when ``search`` declares ``requires_fp64``. + """ + global _x64_warning_emitted + + if not analysis.is_jax: + return + + import jax + + if jax.config.jax_enable_x64: + return + + if type(search).requires_fp64: + raise exc.SearchException( + X64_REQUIRED_MESSAGE.format(search=type(search).__name__) + ) + + if _x64_warning_emitted: + return + + _x64_warning_emitted = True + logger.warning(X64_WARNING) + warnings.warn(X64_WARNING, UserWarning, stacklevel=3) + + +def preflight_kinds(search) -> tuple: + """ + The objective kinds the trace preflight traces for ``search``: its declared + execution kind, plus the matching gradient kind when it declares + ``gradient == "uses"``. + """ + 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) + + +def _preflight_fitness(search, analysis, model): + """ + A throwaway ``Fitness`` in the search's declared convention, 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, + ) + + +def trace_preflight(search, analysis, model): + """ + Trace the search's declared objective kind(s) 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. + + 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): + return + + import jax + import jax.numpy as jnp + + fitness = _preflight_fitness(search, analysis, model) + + vector = jnp.asarray(model.physical_values_from_prior_medians) + + for kind in preflight_kinds(search): + objective = fitness.objective(kind, compile=False) + argument = vector[None, :] if kind in (BATCHED, BATCHED_VALUE_AND_GRAD) else vector + try: + jax.eval_shape(objective, argument) + except Exception as error: + raise exc.SearchException( + PREFLIGHT_MESSAGE.format( + search=type(search).__name__, + kind=kind, + error_type=type(error).__name__, + error=error, + analysis=type(analysis).__name__, + env=PREFLIGHT_ENV, + ) + ) from error + + if _env_flag(PROBE_ENV, default=False): + numerical_probe(search, fitness, vector) + + +def numerical_probe(search, fitness, vector) -> float: + """ + Evaluate the scalar objective once, eagerly, at ``vector`` and warn when it returns + the search's invalid value (a NaN, infinite, over-ceiling or assertion-violating + likelihood at the prior medians). Executes the likelihood; off by default. + + Returns + ------- + The figure of merit at ``vector``. + """ + import numpy as np + + value = float(np.asarray(fitness.objective(SCALAR, compile=False)(vector))) + + invalid = float(type(search).invalid_value) + + if not np.isfinite(value) or value == invalid: + logger.warning( + f"{type(search).__name__} numerical probe: the objective at the prior-median " + f"parameters is {value} (the search's invalid value is {invalid}). The " + f"likelihood traces, but returns an invalid value there; check for NaNs, an " + f"assertion that excludes the prior medians, or a log likelihood over the " + f"configured ceiling." + ) + + return value diff --git a/autofit/non_linear/search/registry.py b/autofit/non_linear/search/registry.py index 2b95f21ad..35a03cd9b 100644 --- a/autofit/non_linear/search/registry.py +++ b/autofit/non_linear/search/registry.py @@ -150,6 +150,7 @@ def _capabilities( objective_target: Tuple[str, str], invalid_value: str, checkpointer: str = "dill", + requires_fp64: bool = False, ) -> Dict[str, Any]: quantity, space = objective_target values = dict( @@ -169,6 +170,7 @@ def _capabilities( objective_target={"quantity": quantity, "space": space}, invalid_value=invalid_value, checkpointer=checkpointer, + requires_fp64=requires_fp64, ) return {name: values[name] for name in CAPABILITY_ATTRIBUTES} diff --git a/docs/design/checkpointing.md b/docs/design/checkpointing.md index 31e766c0f..783a283aa 100644 --- a/docs/design/checkpointing.md +++ b/docs/design/checkpointing.md @@ -33,7 +33,7 @@ afresh, so they are not resumable and have no resume state. | BlackJAX NUTS, SMC | `PickleCheckpointer()` (`search_internal.pickle`, retained) | none | | DynestyStatic, DynestyDynamic | `DillCheckpointer()` | `NativeFileCheckpointer("savestate.save", .restore)` | | Nautilus | `DillCheckpointer(strip_attributes=("pool_l", "pool_s"))` | `NativeFileCheckpointer("checkpoint.hdf5", ...)`; only nautilus can reopen it (inside a `Sampler` built with the run's prior and likelihood), so its loader raises `NotImplementedError` | -| NSS | `DillCheckpointer()` | `NativeFileCheckpointer("nss_checkpoint.pkl", _load_checkpoint)` (written, read and deleted by NSS's own `_fit`) | +| NSS | `DillCheckpointer()` | `NativeFileCheckpointer("nss_checkpoint.pkl", _load_checkpoint)` (written, read and deleted by NSS's `run(ctx)`) | | MultiStart* | `DillCheckpointer()` | the same file | Every existing filename is kept, so output folders written before A3 still load. diff --git a/docs/design/run_ctx.md b/docs/design/run_ctx.md index 1f088d04b..2b30ff440 100644 --- a/docs/design/run_ctx.md +++ b/docs/design/run_ctx.md @@ -149,3 +149,22 @@ pins the following details the table left open; each is flagged for review in th `samples_from_raw` are A3's `autofit/non_linear/samples/adapter.py` (re-exported from `autofit/non_linear/search/fit_context.py`), and `samples_via_internal_from` is implemented once on `NonLinearSearch` over `raw_samples_from` for every search. + +## Revision 2 (phase A3b, 2026-10-08): NSS on the bridge + +The signature and the member list are unchanged. + +- **`NSS` moved onto `run(ctx)`** ahead of A5, because A3b needed it on `Fitness`: it samples + through `ctx.fitness.objective("scalar", compile=False)` (blackjax traces it inside its own + jitted step), resumes from `ctx.resume` (its `NativeFileCheckpointer("nss_checkpoint.pkl")`) + and switches the `Fitness` quick-update cadence off through `fitness_overrides`, since it + fires its own between outer iterations. `make_fitness` is now the only `Fitness` + construction site in `autofit/non_linear/search`, and NSS's resume likelihood sanity check + runs when the bridge builds that `Fitness`, before sampling. +- **The bridge repeats the REQUIRED gate** (`capabilities.check_jax_required`) before it + builds the `Fitness`, so a caller that reaches `_fit` without `fit` gets the shared message + for every `run(ctx)` search, not a tracer error from the backend. +- **The JAX preflight** (`autofit/non_linear/search/preflight.py`) runs in `start_resume_fit` + after the test-mode bypass and the REQUIRED gate, before the context exists: the x64 check + and one `jax.eval_shape` trace of the declared objective kind (and its gradient kind for + `gradient == "uses"`). `run` therefore only ever sees a likelihood that traces. 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 0188df1d1..e3789e3dd 100644 --- a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py +++ b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py @@ -156,20 +156,20 @@ def test__init_iterations_per_quick_update_no_longer_warns(caplog): @requires_nss -def test__load_checkpoint_called_when_file_exists(tmp_path): - """Verify ``_load_checkpoint`` is invoked from ``_fit`` when a checkpoint - file exists at the resolved path. +def test__resume_state_loaded_when_ctx_resume_is_set(tmp_path): + """Verify ``run(ctx)`` resumes from ``ctx.resume`` (NSS's ``nss_checkpoint.pkl`` + resume state) instead of initialising fresh particles. We replace ``_blackjax.nss`` with a mock that fails loudly if ``algo.init`` - fires — i.e. if the fresh-init branch ran. The resume branch must call - ``_load_checkpoint`` first; we patch that helper to return a sentinel - that satisfies the immediate-termination logZ check, so the outer loop - exits before doing any real work. + fires — i.e. if the fresh-init branch ran. The resume branch must load + ``ctx.resume`` first; it returns a sentinel state that satisfies the + immediate-termination logZ check, so the outer loop exits before doing + any real work. """ from types import SimpleNamespace + from unittest.mock import MagicMock - fake_checkpoint = tmp_path / "nss_checkpoint.pkl" - fake_checkpoint.write_bytes(b"placeholder - actual contents replaced by mock") + from autofit.non_linear.paths.null import NullPaths sentinel_state = SimpleNamespace( integrator=SimpleNamespace(logZ=0.0, logZ_live=-100.0), @@ -179,32 +179,32 @@ def test__load_checkpoint_called_when_file_exists(tmp_path): ), ) - # Minimum stub for model + analysis. _fit calls model.prior_count and - # both vector_from_unit_vector + log_prior_list_from_vector inside the - # closure construction (which isn't traced unless one_step fires). + # 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, - instance_from_vector=lambda **kw: SimpleNamespace(), 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], ) - mock_analysis = SimpleNamespace( - log_likelihood_function=lambda instance: 0.0, - is_jax=True, + + 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) - # Force the checkpoint property to return our sentinel path even though - # paths is the default NullPaths. + with patch.object( - type(search), - "_nss_checkpoint_path", - new=fake_checkpoint, - ), patch.object( - nss_search_module, - "_load_checkpoint", - return_value=(sentinel_state, [], jax.random.PRNGKey(0), 17), - ) as mock_load, patch.object( nss_search_module, "_blackjax", ) as mock_bjax, patch.object( @@ -216,7 +216,7 @@ def test__load_checkpoint_called_when_file_exists(tmp_path): loglikelihood=np.zeros(1), ), update_info=SimpleNamespace( - num_steps=np.zeros(1, dtype=int), + num_expansions=np.zeros(1, dtype=int), num_shrink=np.zeros(1, dtype=int), ), ), @@ -226,19 +226,13 @@ def test__load_checkpoint_called_when_file_exists(tmp_path): return_value=jnp.zeros((1, 100)), ): mock_bjax.nss.return_value.init.side_effect = AssertionError( - "algo.init called from resume path — expected _load_checkpoint instead." + "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." ) - # The mocked downstream pipeline (Fitness construction, _NSSInternal - # repackaging) is intentionally not realistic — we only care that the - # resume branch was entered and ``_load_checkpoint`` was called. Catch - # any downstream stub-related failure; the assertion below is the gate. - try: - search._fit(model=mock_model, analysis=mock_analysis) - except (AttributeError, AssertionError, TypeError): - pass + internal = search.run(ctx) - mock_load.assert_called_once_with(fake_checkpoint) + resume.load.assert_called_once_with(paths) + assert isinstance(internal, nss_search_module._NSSInternal) diff --git a/test_autofit/non_linear/search/nest/nss/test_log_likelihood_ceiling.py b/test_autofit/non_linear/search/nest/nss/test_log_likelihood_ceiling.py index 1623b6f73..9efbc83dd 100644 --- a/test_autofit/non_linear/search/nest/nss/test_log_likelihood_ceiling.py +++ b/test_autofit/non_linear/search/nest/nss/test_log_likelihood_ceiling.py @@ -1,26 +1,30 @@ """ -The log-likelihood guards on the closure ``af.NSS`` samples through. +The log-likelihood guards on the objective ``af.NSS`` samples through. -``af.NSS`` does not route sampling through `Fitness.call` — it hands blackjax an inline JAX -closure — so the guards there are a second implementation and need their own test. The closure is -built by `nss_log_likelihood_from`, and these tests call that factory rather than restating its -body, so a change to the sampled path cannot pass a test that only resembles it. +``af.NSS`` hands blackjax its fit's scalar objective (``Fitness.objective("scalar")``, +built by ``NonLinearSearch.make_fitness`` from NSS's declared ``objective_target`` and +``invalid_value``), so it shares ``Fitness.call``'s one guard path. These tests build that +objective exactly as ``NSS.run`` does and check the values NSS's backend sees: before +search-extensibility A3b NSS sampled through a second, inline closure and these tests +pinned that closure instead. -The factory is JAX-only, so the file skips whole when JAX is absent (library policy: unit tests are -the numpy always-green layer). +JAX-only (NSS is ``jax_use='required'``), so the file skips whole when JAX is absent +(library policy: unit tests are the numpy always-green layer). """ import numpy as np import pytest +jnp = pytest.importorskip("jax.numpy") + import autofit as af -from autofit.non_linear.search.nest.nss.search import ( - NSS_INVALID_LOG_LIKELIHOOD, - nss_log_likelihood_from, -) +from autofit.non_linear.search.nest.nss import search as nss_search_module +from autofit.non_linear.search.nest.nss.search import NSS_INVALID_LOG_LIKELIHOOD from test_autofit.non_linear.constant_analysis import ConstantAnalysis -jnp = pytest.importorskip("jax.numpy") +pytestmark = pytest.mark.skipif( + not nss_search_module._HAS_NSS, reason="requires blackjax >= 1.6" +) PARAMETERS = [1.0, 1.0, 1.0] @@ -36,12 +40,21 @@ UNDER_CEILING = 1.0e19 -def _log_likelihood_from(log_likelihood, **kwargs): - return nss_log_likelihood_from( +def _log_likelihood_from(log_likelihood, log_likelihood_ceiling=None): + """ + The objective ``NSS.run`` hands blackjax, built from NSS's ``make_fitness``. + + ``Fitness`` reads the configured ceiling once in ``__init__``; an explicit ceiling is + set on it before the (lazily traced) objective is first called. + """ + fitness = af.NSS().make_fitness( model=af.Model(af.ex.Gaussian), - analysis=ConstantAnalysis(log_likelihood=log_likelihood), - **kwargs, + analysis=ConstantAnalysis(log_likelihood=log_likelihood, use_jax=True), + paths=None, ) + if log_likelihood_ceiling is not None: + fitness.log_likelihood_ceiling = log_likelihood_ceiling + return fitness.objective("scalar", compile=False) def test_nss_closure_rejects_a_log_likelihood_above_the_ceiling(): diff --git a/test_autofit/non_linear/search/nest/nss/test_nss_fitness.py b/test_autofit/non_linear/search/nest/nss/test_nss_fitness.py new file mode 100644 index 000000000..955de30c0 --- /dev/null +++ b/test_autofit/non_linear/search/nest/nss/test_nss_fitness.py @@ -0,0 +1,146 @@ +""" +``af.NSS`` on the shared ``Fitness`` (search-extensibility A3b; survey 02 §3.3, §5.5). + +NSS used to hand blackjax an inline closure of its own, so it missed what ``Fitness`` +gives every other search: + +- (b) a model **assertion** raised ``TracerBoolConversionError`` under ``jax.jit``, + because the closure built instances with the raising check; ``Fitness.call`` applies + the assertions as a traced ``xp.where`` penalty instead; +- (c) its ``Fitness`` was built after sampling, so the resume likelihood sanity check + ran at the end of a resumed run rather than the start. + +NSS is ``jax_use='required'``, so the file skips whole without JAX / blackjax. +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") + +import autofit as af +from autofit.non_linear.fitness import Fitness +from autofit.non_linear.search.nest.nss import search as nss_search_module +from autofit.non_linear.search.nest.nss.search import NSS_INVALID_LOG_LIKELIHOOD + +pytestmark = [ + pytest.mark.skipif( + not nss_search_module._HAS_NSS, reason="requires blackjax >= 1.6" + ), + pytest.mark.filterwarnings("ignore::FutureWarning"), +] + + +class _GaussianJaxAnalysis(af.Analysis): + def __init__(self): + super().__init__(use_jax=True) + + def log_likelihood_function(self, instance): + import jax.numpy as jnp + + return -0.5 * ( + (instance.centre - 50.0) ** 2 + + (instance.normalization - 25.0) ** 2 + + (instance.sigma - 10.0) ** 2 + ) * jnp.ones(()) + + +def _model_with_assertion(): + model = af.Model(af.ex.Gaussian) + model.centre = af.UniformPrior(lower_limit=0.0, upper_limit=100.0) + model.normalization = af.UniformPrior(lower_limit=0.0, upper_limit=50.0) + model.sigma = af.UniformPrior(lower_limit=0.0, upper_limit=20.0) + model.add_assertion(model.centre > model.normalization) + return model + + +def test_nss_objective_applies_assertions_under_jit(): + """ + The scalar objective NSS samples through, jitted the way blackjax's step jits it: + an assertion-violating vector maps to NSS's declared invalid value, without a tracer + error, and a satisfying one keeps its log likelihood. + """ + model = _model_with_assertion() + fitness = af.NSS().make_fitness( + model=model, analysis=_GaussianJaxAnalysis(), paths=None + ) + + objective = jax.jit(fitness.objective("scalar", compile=False)) + + violating = jax.numpy.array([10.0, 25.0, 10.0]) # centre < normalization + satisfying = jax.numpy.array([50.0, 25.0, 10.0]) + + assert float(objective(violating)) == NSS_INVALID_LOG_LIKELIHOOD + assert float(objective(satisfying)) == pytest.approx(0.0) + + +def test_nss_fit_with_assertions_runs_under_jit(monkeypatch): + """ + The whole NSS loop (``jax.jit`` of blackjax's step) on a model with an assertion: + before A3b this raised ``TracerBoolConversionError`` (survey 02 §3.3(b)). The + maximum-likelihood sample satisfies the assertion. + """ + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + + model = _model_with_assertion() + + search = af.NSS(n_live=50, num_delete=10, num_mcmc_steps=3) + + result = search.fit(model=model, analysis=_GaussianJaxAnalysis()) + + instance = result.samples.max_log_likelihood() + assert instance.centre > instance.normalization + + log_likelihoods = np.asarray(result.samples.log_likelihood_list) + assert np.max(log_likelihoods) > NSS_INVALID_LOG_LIKELIHOOD + + +def test_nss_fitness_is_built_before_the_backend_runs(monkeypatch): + """ + The resume likelihood sanity check (``Fitness.check_log_likelihood``, run when a + ``Fitness`` with paths is built) happens before sampling, not after it (survey 02 + §3.3(c)). + """ + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + calls = [] + + original = Fitness.check_log_likelihood + + def recording_check(self, fitness): + calls.append("check") + return original(self, fitness) + + class _Stop(Exception): + pass + + def recording_run(self, ctx): + calls.append("run") + raise _Stop + + monkeypatch.setattr(Fitness, "check_log_likelihood", recording_check) + monkeypatch.setattr(af.NSS, "run", recording_run) + + with pytest.raises(_Stop): + af.NSS().fit(model=_model_with_assertion(), analysis=_GaussianJaxAnalysis()) + + # The preflight's throwaway Fitness has no paths, so it never runs the check. + assert calls == ["check", "run"] + + +def test_nss_switches_the_fitness_quick_update_cadence_off(): + """ + NSS fires its own quick update between outer iterations, so the ``Fitness`` it + samples through carries no cadence (no second background worker). + """ + search = af.NSS(iterations_per_quick_update=5) + + fitness = search.make_fitness( + model=_model_with_assertion(), + analysis=_GaussianJaxAnalysis(), + paths=None, + **search.fitness_overrides(_GaussianJaxAnalysis()), + ) + + assert fitness.iterations_per_quick_update is None + assert search.iterations_per_quick_update == 5 diff --git a/test_autofit/non_linear/search/nest/test_dynesty_single_core_sentinel.py b/test_autofit/non_linear/search/nest/test_dynesty_single_core_sentinel.py new file mode 100644 index 000000000..9cec9945d --- /dev/null +++ b/test_autofit/non_linear/search/nest/test_dynesty_single_core_sentinel.py @@ -0,0 +1,147 @@ +""" +Dynesty's single-core control flow (search-extensibility A3b; survey 01 §8, survey 02 +§3.7). + +``AbstractDynesty._fit`` used ``raise RuntimeError`` inside a ``try`` to select its +single-core path, so its ``except RuntimeError`` also swallowed any ``RuntimeError`` +raised while a pooled run executed (a genuine dynesty error, or JAX's +``XlaRuntimeError``, which subclasses it) and silently restarted the run single-core. +The control flow is now the private ``_SingleCoreRun`` sentinel: only a failure to +*create* the pool still falls back. + +The backend is stubbed (``search_internal_from`` / ``run_search_internal``), so no +sampler runs. +""" + +import pytest + +import autofit as af +from autofit.non_linear.search.nest.dynesty.search import abstract as dynesty_abstract + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +class XlaRuntimeErrorLike(RuntimeError): + """ + Stands in for ``jaxlib.xla_extension.XlaRuntimeError``, a ``RuntimeError`` subclass. + """ + + +class _Analysis(af.Analysis): + def log_likelihood_function(self, instance): + return 0.0 + + +class _RecordingPool: + """ + A context-manager pool whose ``__enter__`` can be made to fail like a platform + without working multiprocessing. + """ + + fail_on_enter = False + entered = 0 + + def __init__(self, **kwargs): + pass + + def __enter__(self): + if self.fail_on_enter: + raise RuntimeError("multiprocessing is not available here") + type(self).entered += 1 + return self + + def __exit__(self, *args): + return False + + +def _search(monkeypatch, pool_cls, run_search_internal): + search = af.DynestyStatic(number_of_cores=2) + + pools = [] + + def search_internal_from(model, fitness, checkpoint_exists, pool, queue_size): + pools.append(pool) + return object() + + monkeypatch.setattr(dynesty_abstract, "_fork_pool_cls", lambda: pool_cls) + monkeypatch.setattr(search, "search_internal_from", search_internal_from) + monkeypatch.setattr(search, "run_search_internal", run_search_internal) + return search, pools + + +def test_runtime_error_inside_a_pooled_run_propagates(monkeypatch): + calls = [] + + def run_search_internal(search_internal): + calls.append(search_internal) + raise XlaRuntimeErrorLike("RESOURCE_EXHAUSTED: out of memory") + + class Pool(_RecordingPool): + entered = 0 + + search, pools = _search(monkeypatch, Pool, run_search_internal) + + with pytest.raises(XlaRuntimeErrorLike, match="RESOURCE_EXHAUSTED"): + search._fit(model=af.Model(af.ex.Gaussian), analysis=_Analysis()) + + # One pooled attempt, no silent single-core rerun. + assert len(calls) == 1 + assert len(pools) == 1 and pools[0] is not None + assert Pool.entered == 1 + + +def test_pool_creation_failure_still_falls_back_to_single_core(monkeypatch, caplog): + def run_search_internal(search_internal): + return True + + class Pool(_RecordingPool): + fail_on_enter = True + + search, pools = _search(monkeypatch, Pool, run_search_internal) + + with caplog.at_level("INFO"): + search._fit(model=af.Model(af.ex.Gaussian), analysis=_Analysis()) + + assert pools == [None] + assert "multiprocessing pool could not be created" in caplog.text + assert "multiprocessing is not available here" in caplog.text + + +def test_single_core_is_selected_without_a_runtime_error(monkeypatch): + def run_search_internal(search_internal): + return True + + def pool_must_not_be_built(): + raise AssertionError("a single-core fit built a pool") + + search = af.DynestyStatic(number_of_cores=1) + pools = [] + + monkeypatch.setattr(dynesty_abstract, "_fork_pool_cls", pool_must_not_be_built) + monkeypatch.setattr( + search, + "search_internal_from", + lambda model, fitness, checkpoint_exists, pool, queue_size: pools.append(pool), + ) + monkeypatch.setattr(search, "run_search_internal", run_search_internal) + + search._fit(model=af.Model(af.ex.Gaussian), analysis=_Analysis()) + + assert pools == [None] + assert not issubclass(dynesty_abstract._SingleCoreRun, RuntimeError) + + +def test_runtime_error_in_a_single_core_run_propagates(monkeypatch): + def run_search_internal(search_internal): + raise XlaRuntimeErrorLike("INTERNAL") + + search = af.DynestyStatic(number_of_cores=1) + monkeypatch.setattr( + search, + "search_internal_from", + lambda model, fitness, checkpoint_exists, pool, queue_size: object(), + ) + monkeypatch.setattr(search, "run_search_internal", run_search_internal) + + with pytest.raises(XlaRuntimeErrorLike): + search._fit(model=af.Model(af.ex.Gaussian), analysis=_Analysis()) diff --git a/test_autofit/non_linear/search/test_make_fitness.py b/test_autofit/non_linear/search/test_make_fitness.py index 44c1d3d8f..2ddf44b96 100644 --- a/test_autofit/non_linear/search/test_make_fitness.py +++ b/test_autofit/non_linear/search/test_make_fitness.py @@ -7,7 +7,7 @@ forgot it silently never fired quick updates (PyAutoFit#1434). With every search building its `Fitness` through `make_fitness`, which always forwards the cadence, that scan has nothing left to guard; what remains is that `make_fitness` forwards it, and that -no search builds a `Fitness` any other way (NSS keeps its own path until phase A3b). +no search builds a `Fitness` any other way (NSS, the last exception, moved onto it in A3b). """ import ast from pathlib import Path @@ -89,8 +89,9 @@ def _fitness_constructions(): yield path.relative_to(SEARCH_ROOT).as_posix() -def test_only_make_fitness_and_nss_construct_a_fitness(): - assert sorted(_fitness_constructions()) == [ - "abstract_search.py", - "nest/nss/search.py", - ] +def test_only_make_fitness_constructs_a_fitness(): + """ + NSS, the last search that built its own ``Fitness``, moved onto ``run(ctx)`` in + search-extensibility A3b; ``make_fitness`` is now the only construction site. + """ + assert sorted(_fitness_constructions()) == ["abstract_search.py"] diff --git a/test_autofit/non_linear/search/test_preflight.py b/test_autofit/non_linear/search/test_preflight.py new file mode 100644 index 000000000..08c98495c --- /dev/null +++ b/test_autofit/non_linear/search/test_preflight.py @@ -0,0 +1,299 @@ +""" +The JAX preflight ``NonLinearSearch.fit`` runs before a backend builds any state +(search-extensibility A3b; survey 02 §5.3, ``autofit.non_linear.search.preflight``). + +- the trace preflight: one ``jax.eval_shape`` of the declared objective kind (and of its + gradient kind for ``gradient == "uses"``) at the prior medians, raising a + ``SearchException`` that names the search with the tracer error chained; +- the optional numerical probe (``PYAUTO_JAX_PREFLIGHT_PROBE=1``), off by default; +- the x64 check: warn once when a JAX analysis runs with ``jax_enable_x64`` off, raise + for a search declaring ``requires_fp64``; +- placement: after the test-mode bypass (``PYAUTO_TEST_MODE>=2`` never runs it) and after + the ``jax_use='required'`` gate. + +No backend executes: every fit here either raises before ``_fit`` or bypasses it. +""" + +import logging + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") + +import autofit as af +from autofit import exc +from autofit.non_linear.search import capabilities as cap +from autofit.non_linear.search import preflight + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +class NumpyCastAnalysis(af.Analysis): + """ + A ``use_jax=True`` analysis whose likelihood casts a parameter with + ``np.asarray``: it cannot be traced. + """ + + def __init__(self): + super().__init__(use_jax=True) + + def log_likelihood_function(self, instance): + centre = np.asarray(instance.centre) + return -0.5 * float((centre - 50.0) ** 2) + + +class TraceableAnalysis(af.Analysis): + def __init__(self, value=None): + super().__init__(use_jax=True) + self.value = value + + def log_likelihood_function(self, instance): + import jax.numpy as jnp + + if self.value is not None: + return jnp.asarray(self.value) + return -0.5 * (instance.centre - 50.0) ** 2 + + +def _model(): + return af.Model(af.ex.Gaussian) + + +def _backend_must_not_run(*args, **kwargs): + raise AssertionError("the backend ran before the preflight") + + +def test_preflight_raises_on_an_np_asarray_likelihood_with_the_error_chained( + monkeypatch, +): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + search = af.DynestyStatic() + monkeypatch.setattr(search, "_fit", _backend_must_not_run) + + with pytest.raises(exc.SearchException) as error: + search.fit(model=_model(), analysis=NumpyCastAnalysis()) + + message = str(error.value) + assert "DynestyStatic" in message + assert "'scalar'" in message + assert "TracerArrayConversionError" in message + assert preflight.PREFLIGHT_ENV in message + + cause = error.value.__cause__ + assert cause is not None + assert type(cause).__name__ == "TracerArrayConversionError" + + +def test_preflight_traces_the_batched_kind_of_a_batched_search(): + with pytest.raises(exc.SearchException, match="'batched'") as error: + preflight.trace_preflight(af.Nautilus(), NumpyCastAnalysis(), _model()) + + assert "Nautilus" in str(error.value) + assert error.value.__cause__ is not None + + +def test_preflight_kinds_follow_the_declared_capabilities(): + assert preflight.preflight_kinds(af.DynestyStatic()) == ("scalar",) + assert preflight.preflight_kinds(af.Nautilus()) == ("batched",) + + class GradientSearch(af.DynestyStatic): + gradient = cap.Gradient.USES + + class BatchedGradientSearch(af.DynestyStatic): + gradient = cap.Gradient.USES + batched = True + + assert preflight.preflight_kinds(GradientSearch()) == ("scalar", "value_and_grad") + assert preflight.preflight_kinds(BatchedGradientSearch()) == ( + "batched", + "batched_value_and_grad", + ) + + +def test_preflight_traces_the_gradient_of_a_gradient_search(monkeypatch): + """ + A gradient search's gradient kind is traced too: the trace sees the gradient's own + failure (here a function JAX can evaluate but not differentiate). + """ + import jax.numpy as jnp + + @jax.custom_jvp + def no_derivative(x): + return x + + @no_derivative.defjvp + def _no_derivative_jvp(primals, tangents): + raise NotImplementedError("this likelihood has no derivative") + + class NonDifferentiableAnalysis(af.Analysis): + def __init__(self): + super().__init__(use_jax=True) + + def log_likelihood_function(self, instance): + return -0.5 * no_derivative(instance.centre - 50.0) ** 2 * jnp.ones(()) + + class GradientSearch(af.DynestyStatic): + gradient = cap.Gradient.USES + + preflight.trace_preflight(af.DynestyStatic(), NonDifferentiableAnalysis(), _model()) + + with pytest.raises(exc.SearchException, match="'value_and_grad'") as error: + preflight.trace_preflight(GradientSearch(), NonDifferentiableAnalysis(), _model()) + + assert isinstance(error.value.__cause__, NotImplementedError) + + +def test_preflight_passes_a_traceable_likelihood_and_evaluates_nothing(monkeypatch): + calls = [] + + class CountingAnalysis(TraceableAnalysis): + def log_likelihood_function(self, instance): + calls.append(type(instance.centre).__name__) + return super().log_likelihood_function(instance) + + preflight.trace_preflight(af.DynestyStatic(), CountingAnalysis(), _model()) + + # One trace: the likelihood saw tracers, never concrete values. + assert len(calls) == 1 + assert "Tracer" in calls[0] or "ShapedArray" in calls[0] + + +def test_preflight_is_skipped_by_its_environment_flag(monkeypatch): + monkeypatch.setenv(preflight.PREFLIGHT_ENV, "0") + preflight.trace_preflight(af.DynestyStatic(), NumpyCastAnalysis(), _model()) + + +def test_preflight_is_a_no_op_for_a_numpy_analysis(): + class NumpyAnalysis(af.Analysis): + def log_likelihood_function(self, instance): + return float(np.asarray(instance.centre)) + + preflight.trace_preflight(af.DynestyStatic(), NumpyAnalysis(), _model()) + + +def test_preflight_is_bypassed_by_test_mode_2(monkeypatch): + """ + ``PYAUTO_TEST_MODE=2`` returns before the preflight, so even a non-traceable JAX + likelihood completes a smoke run (the bypass evaluates it eagerly). + """ + monkeypatch.setenv("PYAUTO_TEST_MODE", "2") + + result = af.DynestyStatic().fit(model=_model(), analysis=NumpyCastAnalysis()) + + assert isinstance(result, af.Result) + + +def test_required_gate_runs_before_the_preflight(monkeypatch): + """ + A ``jax_use='required'`` search given a numpy analysis gets the shared REQUIRED + message, not a preflight error. + """ + pytest.importorskip("optax") + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + def preflight_must_not_run(*args, **kwargs): + raise AssertionError("the preflight ran before the REQUIRED gate") + + monkeypatch.setattr(preflight, "trace_preflight", preflight_must_not_run) + + class NumpyAnalysis(af.Analysis): + def log_likelihood_function(self, instance): + return 0.0 + + with pytest.raises(exc.SearchException) as error: + af.MultiStartAdam().fit(model=_model(), analysis=NumpyAnalysis()) + + assert str(error.value) == cap.JAX_REQUIRED_MESSAGE.format(search="MultiStartAdam") + + +def test_numerical_probe_is_off_by_default(monkeypatch, caplog): + monkeypatch.delenv(preflight.PROBE_ENV, raising=False) + + with caplog.at_level(logging.WARNING, logger=preflight.__name__): + preflight.trace_preflight( + af.DynestyStatic(), TraceableAnalysis(value=np.nan), _model() + ) + + assert "numerical probe" not in caplog.text + + +def test_numerical_probe_warns_on_an_invalid_value_at_the_prior_medians( + monkeypatch, caplog +): + monkeypatch.setenv(preflight.PROBE_ENV, "1") + + with caplog.at_level(logging.WARNING, logger=preflight.__name__): + preflight.trace_preflight( + af.DynestyStatic(), TraceableAnalysis(value=np.nan), _model() + ) + + assert "DynestyStatic numerical probe" in caplog.text + + +def test_numerical_probe_is_quiet_on_a_valid_value(monkeypatch, caplog): + monkeypatch.setenv(preflight.PROBE_ENV, "1") + + with caplog.at_level(logging.WARNING, logger=preflight.__name__): + preflight.trace_preflight(af.DynestyStatic(), TraceableAnalysis(), _model()) + + assert "numerical probe" not in caplog.text + + +# --- x64 ------------------------------------------------------------------------------ + + +@pytest.fixture +def float32(monkeypatch): + """ + Run the test with ``jax_enable_x64`` off, restoring it (and the once-per-process + warning flag) afterwards. + """ + previous = bool(jax.config.jax_enable_x64) + monkeypatch.setattr(preflight, "_x64_warning_emitted", False) + jax.config.update("jax_enable_x64", False) + try: + yield + finally: + jax.config.update("jax_enable_x64", previous) + + +def test_fp32_warns_once(float32, caplog): + analysis = TraceableAnalysis() + + with pytest.warns(UserWarning, match="32-bit precision") as record: + with caplog.at_level(logging.WARNING, logger=preflight.__name__): + preflight.check_x64(af.DynestyStatic(), analysis) + preflight.check_x64(af.DynestyStatic(), analysis) + + assert len([w for w in record if "32-bit" in str(w.message)]) == 1 + assert caplog.text.count("32-bit precision") == 1 + + +def test_fp32_raises_for_a_search_declaring_requires_fp64(float32): + class Fp64Search(af.DynestyStatic): + requires_fp64 = True + + with pytest.raises(exc.SearchException, match="Fp64Search requires float64"): + preflight.check_x64(Fp64Search(), TraceableAnalysis()) + + +def test_fp32_check_is_silent_for_numpy_and_for_x64(monkeypatch, recwarn): + monkeypatch.setattr(preflight, "_x64_warning_emitted", False) + + class NumpyAnalysis(af.Analysis): + def log_likelihood_function(self, instance): + return 0.0 + + preflight.check_x64(af.DynestyStatic(), NumpyAnalysis()) + + if jax.config.jax_enable_x64: + preflight.check_x64(af.DynestyStatic(), TraceableAnalysis()) + + assert not [w for w in recwarn if "32-bit" in str(w.message)] + + +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())