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

.github/workflows/docs.yml

@@ -29,3 +29,18 @@ jobs:
     with:
       package: autofit
     secrets: inherit
+
+  # The search pages (docs/api/searches.rst, docs/searches/index.rst and
+  # docs/searches/citations.rst) are generated from the search manifest
+  # (python -m autofit search-manifest --json); fail when they are stale.
+  generated-search-docs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install autofit
+        run: pip install -e .
+      - name: Check the generated search docs
+        run: python docs/_generate_searches.py --check

autofit/__main__.py

@@ -0,0 +1,58 @@
+"""
+Command-line entry point: ``python -m autofit <command>``.
+
+Commands
+--------
+search-manifest [--json]
+    Print the declarative search registry. ``--json`` emits the versioned manifest
+    (``search-manifest@1``), the only cross-repo format for search facts; without it a
+    plain-text table is printed.
+"""
+
+import argparse
+import json
+import sys
+
+
+def _search_manifest(arguments) -> int:
+    from autofit.non_linear.search.registry import manifest
+
+    data = manifest()
+
+    if arguments.json:
+        json.dump(data, sys.stdout, indent=2)
+        sys.stdout.write("\n")
+        return 0
+
+    print(f"{data['schema']} (autofit {data['autofit_version']})")
+    for search in data["searches"]:
+        capabilities = search["capabilities"]
+        print(
+            f"{search['name']:<20} {search['family']:<5} "
+            f"jax={capabilities['jax_use']:<9} gradient={capabilities['gradient']:<5} "
+            f"posterior={capabilities['posterior_kind']:<9} "
+            f"evidence={capabilities['produces_evidence']!s:<5} "
+            f"status={capabilities['status']}"
+        )
+    return 0
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(prog="python -m autofit")
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    search_manifest = commands.add_parser(
+        "search-manifest",
+        help="Print the declarative search registry (search-manifest@1).",
+    )
+    search_manifest.add_argument(
+        "--json", action="store_true", help="Emit the versioned JSON manifest."
+    )
+    search_manifest.set_defaults(handler=_search_manifest)
+
+    arguments = parser.parse_args(argv)
+    return arguments.handler(arguments)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

autofit/graphical/declarative/collection.py

@@ -1,4 +1,4 @@
-from typing import Union, Optional
+from typing import List, Union, Optional
 
 from autofit.graphical.declarative.factor.hierarchical import HierarchicalFactor
 
@@ -19,7 +19,7 @@ def __init__(
         *model_factors: Union[AbstractDeclarativeFactor, HierarchicalFactor],
         name=None,
         include_prior_factors=True,
-        use_jax : bool = False
+        use_jax: Optional[bool] = None,
     ):
         """
         A collection of factors that describe models, which can be
@@ -31,6 +31,15 @@ def __init__(
         ----------
         model_factors
             Factors which are hierarchical or associated with a specific analysis
+        use_jax
+            Whether the whole graph is evaluated on JAX when it is fitted as one
+            analysis (``search.fit(model=graph.global_prior_model, analysis=graph)``).
+            ``None`` (the default) derives it from the factors: the graph is JAX when
+            every factor is. An explicit bool is honoured, but whole-graph fitting
+            requires every factor to agree with it and raises ``SearchException``
+            naming the factors that do not (``check_backend_agreement``). Per-factor
+            expectation propagation (``graph.optimise``) fits each factor with its own
+            analysis, so mixed numpy/JAX factors are allowed there.
         """
         super().__init__(
             include_prior_factors=include_prior_factors,
@@ -42,7 +51,7 @@ def __init__(
     def tree_flatten(self):
         return (
             (self._model_factors,),
-            (self._name, self._include_prior_factors),
+            (self._name, self._include_prior_factors, self._explicit_use_jax),
         )
 
     @classmethod
@@ -51,8 +60,103 @@ def tree_unflatten(cls, aux_data, children):
             *children[0],
             name=aux_data[0],
             include_prior_factors=aux_data[1],
+            use_jax=aux_data[2],
         )
 
+    def __setstate__(self, state):
+        """
+        Restore a pickled graph, migrating pickles written before ``use_jax`` could be
+        derived from the factors.
+
+        Those stored the graph's backend as a plain ``_use_jax`` attribute, which the
+        ``_use_jax`` property below now shadows; it is moved to ``_explicit_use_jax``
+        so the graph keeps the backend it was built with (then always an explicit
+        bool, ``False`` by default) and ``tree_flatten`` finds the attribute.
+        """
+        state = dict(state)
+        if "_explicit_use_jax" not in state:
+            state["_explicit_use_jax"] = state.pop("_use_jax", None)
+        self.__dict__.update(state)
+
+    @property
+    def _use_jax(self) -> bool:
+        """
+        The graph's backend: the explicit ``use_jax`` when one was given, otherwise
+        ``True`` exactly when every factor is JAX (``is_jax`` reads this).
+
+        The factors are the flattened ones (a ``HierarchicalFactor`` contributes the
+        factors it generates), the same set ``check_backend_agreement`` reads, so a
+        derived backend always agrees with its own factors.
+        """
+        explicit = self.__dict__.get("_explicit_use_jax")
+        if explicit is not None:
+            return bool(explicit)
+        flags = [factor.is_jax for factor in self._flat_factors()]
+        return bool(flags) and all(flags)
+
+    @_use_jax.setter
+    def _use_jax(self, value: Optional[bool]):
+        # Set by `Analysis.__init__`: `None` means "derive from the factors".
+        self._explicit_use_jax = value
+
+    def _flat_factors(self):
+        model_factors = list()
+        for model_factor in self.__dict__.get("_model_factors", []):
+            if isinstance(model_factor, HierarchicalFactor):
+                model_factors.extend(model_factor.factors)
+            else:
+                model_factors.append(model_factor)
+        return model_factors
+
+    def factors_disagreeing_on_backend(self) -> List[str]:
+        """
+        The names of the factors whose ``is_jax`` differs from the graph's.
+        """
+        is_jax = self.is_jax
+        return [
+            factor.name for factor in self._flat_factors() if factor.is_jax != is_jax
+        ]
+
+    def check_backend_agreement(self):
+        """
+        Raise ``SearchException`` unless every factor agrees with the graph's backend.
+
+        Whole-graph fitting evaluates every factor inside one objective, so a JAX graph
+        traces numpy factors (which fail on tracers) and a numpy graph never compiles a
+        JAX factor. ``NonLinearSearch`` calls this before fitting a graph as one
+        analysis; per-factor expectation propagation never does.
+        """
+        disagreeing = self.factors_disagreeing_on_backend()
+        if disagreeing:
+            from autofit import exc
+
+            raise exc.SearchException(
+                f"The FactorGraphModel {self.name} is fitted as one analysis with "
+                f"is_jax={self.is_jax}, but these factors have "
+                f"is_jax={not self.is_jax}: {', '.join(disagreeing)}. Whole-graph "
+                f"fitting requires every factor to use the same backend: build every "
+                f"factor's analysis with the same use_jax (and HierarchicalFactor with "
+                f"the same use_jax), or omit FactorGraphModel(use_jax=...) so the graph "
+                f"derives it from its factors. Per-factor expectation propagation "
+                f"(graph.optimise) allows mixed factors."
+            )
+
+    @property
+    def gradient_mode(self) -> str:
+        """
+        The factors' common ``gradient_mode``; ``"reverse"`` (the ``Analysis`` default)
+        when the factors declare different modes. Both modes compute the same gradient,
+        so a mixed graph falls back to the mode whose memory does not scale with the
+        number of free parameters.
+        """
+        modes = {
+            getattr(factor, "gradient_mode", "reverse")
+            for factor in self._flat_factors()
+        }
+        if len(modes) == 1:
+            return modes.pop()
+        return "reverse"
+
     @property
     def prior_model(self):
         """
@@ -144,13 +248,7 @@ def _shared_state_from(self, instance: ModelInstance):
 
     @property
     def model_factors(self):
-        model_factors = list()
-        for model_factor in self._model_factors:
-            if isinstance(model_factor, HierarchicalFactor):
-                model_factors.extend(model_factor.factors)
-            else:
-                model_factors.append(model_factor)
-        return model_factors
+        return self._flat_factors()
 
     def make_result(
         self,

autofit/graphical/declarative/factor/analysis.py

@@ -86,6 +86,27 @@ def __init__(
             name=name,
         )
 
+    @property
+    def is_jax(self) -> bool:
+        """
+        Whether this factor's analysis is JAX (``Analysis.is_jax`` of the wrapped
+        analysis). Per-factor expectation propagation fits each factor with its own
+        analysis, so mixed numpy/JAX factors are allowed there; whole-graph fitting of a
+        ``FactorGraphModel`` requires every factor to agree.
+        """
+        return self.analysis.is_jax
+
+    @property
+    def _use_jax(self) -> bool:
+        return self.analysis.is_jax
+
+    @property
+    def gradient_mode(self) -> str:
+        """
+        The wrapped analysis's declared gradient mode.
+        """
+        return self.analysis.gradient_mode
+
     def tree_flatten(self):
         return (
             (self.prior_model,),

autofit/graphical/declarative/factor/hierarchical.py

@@ -78,6 +78,14 @@ def __init__(
         self._name = name or namer(self.__class__.__name__)
         self._factors = list()
         self.optimiser = optimiser
+
+        # The same switch `Analysis.__init__` honours, so a hierarchical factor in a
+        # graph agrees with its analysis factors when JAX is disabled for a smoke run.
+        import os
+
+        if os.environ.get("PYAUTO_DISABLE_JAX") == "1":
+            use_jax = False
+
         self._use_jax = use_jax
 
     @property

autofit/non_linear/analysis/analysis.py

@@ -151,9 +151,26 @@ def method(*args, **kwargs):
 
         return method
 
+    @property
+    def is_jax(self) -> bool:
+        """
+        Whether this analysis evaluates its likelihood on the JAX backend.
+
+        The single read-only probe every search, ``Fitness`` and the latent machinery use
+        to ask "is this analysis JAX?". ``use_jax`` passed to ``__init__`` is the user's
+        choice; ``is_jax`` is the resolved value after ``PYAUTO_DISABLE_JAX=1`` and the
+        no-JAX-installed fallback have been applied. An analysis that never called
+        ``Analysis.__init__`` is numpy.
+
+        Wrappers forward it: ``ModelAnalysis`` takes its wrapped analysis's value,
+        ``AnalysisFactor`` reads its analysis and ``FactorGraphModel`` derives it from its
+        factors.
+        """
+        return bool(getattr(self, "_use_jax", False))
+
     @property
     def _xp(self):
-        if self._use_jax:
+        if self.is_jax:
             import jax.numpy as jnp
             return jnp
         return np
@@ -356,7 +373,7 @@ def supports_background_update(self) -> bool:
     @property
     def supports_jax_visualization(self) -> bool:
         """Whether the visualizer can work directly with JAX arrays."""
-        return self._use_jax
+        return self.is_jax
 
     def perform_quick_update(self, paths, instance):
         raise NotImplementedError
@@ -388,7 +405,7 @@ def batched_memory_bytes(self, model, batch_size: int, gradient: bool = False):
         The projected bytes, or ``None`` when this analysis is not on the JAX
         path (nothing to measure).
         """
-        if not self._use_jax:
+        if not self.is_jax:
             return None
 
         import jax
@@ -441,7 +458,7 @@ def print_vram_use(self, model, batch_size : int, gradient: bool = False) -> str
         if skip_fit_output():
             return
 
-        if not self._use_jax:
+        if not self.is_jax:
             print("use_jax=False for this analysis, therefore does not use GPU and VRAM use cannot be profiled.")
             return
 

autofit/non_linear/analysis/latent.py

@@ -156,7 +156,7 @@ def _record_failure():
             if failures["first_traceback"] is None:
                 failures["first_traceback"] = traceback.format_exc()
 
-        if analysis._use_jax:
+        if analysis.is_jax:
             import jax
             import jax.numpy as jnp
             if batch_mode == "vmap":
@@ -248,7 +248,7 @@ def batched_compute_latent(x):
             # batched JAX call on this chunk
             latent_values_batch = batched_compute_latent(batch)
 
-            if analysis._use_jax:
+            if analysis.is_jax:
                 import jax.numpy as jnp
                 latent_values_batch = jnp.stack(latent_values_batch, axis=-1)  # (batch, n_latents)
 

autofit/non_linear/analysis/model_analysis.py

@@ -6,20 +6,42 @@
 
 
 class ModelAnalysis(Analysis):
-    def __init__(self, analysis: Analysis, model: AbstractPriorModel, use_jax : bool = False):
+    def __init__(
+        self,
+        analysis: Analysis,
+        model: AbstractPriorModel,
+        use_jax: Optional[bool] = None,
+    ):
         """
         Comprises a model and an analysis that can be applied to instances of that model.
 
         Parameters
         ----------
         analysis
+            The wrapped analysis.
         model
+            The model whose instances the wrapped analysis is evaluated on.
+        use_jax
+            ``None`` (the default) forwards the wrapped analysis's ``is_jax``, so the
+            wrapper and the analysis it wraps can never disagree about the backend. An
+            explicit bool is honoured as before.
         """
+        if use_jax is None:
+            use_jax = analysis.is_jax
+
         super().__init__(use_jax=use_jax)
 
         self.analysis = analysis
         self.model = model
 
+    @property
+    def gradient_mode(self) -> str:
+        """
+        The wrapped analysis's declared gradient mode (``Analysis.gradient_mode``), so a
+        gradient search fitting the wrapper differentiates the way the analysis asked.
+        """
+        return self.analysis.gradient_mode
+
     def __getattr__(self, item):
         if item in ("__getstate__", "__setstate__"):
             raise AttributeError(item)