Merge main into feature/rst-to-myst-md (port Colab tag bumps to .md)

Catches the branch up to main. The 2026.5.1.1 and 2026.5.1.4 release
commits bumped the Colab URL tag in docs/index.rst on main; that file
was deleted in the conversion, so the modify/delete conflict is
resolved by keeping the deletion and porting the bump
(2026.4.13.6 → 2026.5.1.4) to the MyST docs/index.md sibling.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 2 people

README.rst

@@ -2,7 +2,7 @@ PyAutoFit: Classy Probabilistic Programming
 ===========================================
 
 .. |colab| image:: https://colab.research.google.com/assets/colab-badge.svg
-   :target: https://colab.research.google.com/github/PyAutoLabs/autofit_workspace/blob/2026.4.13.6/start_here.ipynb
+   :target: https://colab.research.google.com/github/PyAutoLabs/autofit_workspace/blob/2026.5.1.4/start_here.ipynb
 
 .. |RTD| image:: https://readthedocs.org/projects/pyautofit/badge/?version=latest
     :target: https://pyautofit.readthedocs.io/en/latest/?badge=latest
@@ -33,7 +33,7 @@ PyAutoFit: Classy Probabilistic Programming
 
 `Installation Guide <https://pyautofit.readthedocs.io/en/latest/installation/overview.html>`_ |
 `readthedocs <https://pyautofit.readthedocs.io/en/latest/index.html>`_ |
-`Introduction on Colab <https://colab.research.google.com/github/PyAutoLabs/autofit_workspace/blob/2026.4.13.6/notebooks/overview/overview_1_the_basics.ipynb>`_ |
+`Introduction on Colab <https://colab.research.google.com/github/PyAutoLabs/autofit_workspace/blob/2026.5.1.4/notebooks/overview/overview_1_the_basics.ipynb>`_ |
 `HowToFit <https://github.com/PyAutoLabs/HowToFit>`_
 
 **PyAutoFit** is a Python based probabilistic programming language for model fitting and Bayesian inference
@@ -55,7 +55,7 @@ The following links are useful for new starters:
 
 - `The PyAutoFit readthedocs <https://pyautofit.readthedocs.io/en/latest>`_, which includes an `installation guide <https://pyautofit.readthedocs.io/en/latest/installation/overview.html>`_ and an overview of **PyAutoFit**'s core features.
 
-- `The introduction Jupyter Notebook on Colab <https://colab.research.google.com/github/PyAutoLabs/autofit_workspace/blob/2026.4.13.6/notebooks/overview/overview_1_the_basics.ipynb>`_, where you can try **PyAutoFit** in a web browser (without installation).
+- `The introduction Jupyter Notebook on Colab <https://colab.research.google.com/github/PyAutoLabs/autofit_workspace/blob/2026.5.1.4/notebooks/overview/overview_1_the_basics.ipynb>`_, where you can try **PyAutoFit** in a web browser (without installation).
 
 - `The autofit_workspace GitHub repository <https://github.com/Jammy2211/autofit_workspace>`_, which includes example scripts demonstrating **PyAutoFit**'s features.
 

autofit/__init__.py

@@ -15,6 +15,7 @@
 from .database.aggregator.aggregator import GridSearchAggregator
 from .graphical.expectation_propagation.history import EPHistory
 from .graphical.declarative.factor.analysis import AnalysisFactor
+from .graphical.declarative.factor.analysis import EPAnalysisFactor
 from .graphical.declarative.collection import FactorGraphModel
 from .graphical.declarative.factor.hierarchical import HierarchicalFactor
 from .graphical.laplace import LaplaceOptimiser
@@ -138,7 +139,7 @@ def save_abc(pickler, obj):
     pickle._Pickler.save_type(pickler, obj)
 
 
-__version__ = "2026.4.13.6"
+__version__ = "2026.5.1.4"
 
 from autoconf import check_version
 

autofit/graphical/__init__.py

@@ -1,7 +1,7 @@
 from . import utils
 from .declarative.abstract import PriorFactor
 from .declarative.collection import FactorGraphModel
-from .declarative.factor.analysis import AnalysisFactor
+from .declarative.factor.analysis import AnalysisFactor, EPAnalysisFactor
 from .declarative.factor.hierarchical import _HierarchicalFactor, HierarchicalFactor
 from .expectation_propagation.ep_mean_field import EPMeanField
 from .expectation_propagation.optimiser import EPOptimiser

autofit/graphical/declarative/factor/analysis.py

@@ -187,3 +187,52 @@ def save_results(self, paths: AbstractPaths, result):
 
     def log_likelihood_function(self, instance: ModelInstance) -> float:
         return self.analysis.log_likelihood_function(instance)
+
+
+class EPAnalysisFactor(AnalysisFactor):
+    """
+    An ``AnalysisFactor`` that exposes the EP cavity distribution to its
+    ``Analysis`` on each likelihood evaluation.
+
+    On every iteration of the EP optimiser, the cavity distribution
+    ``q⁻ᵃ`` — the product of the posterior approximations from all
+    *other* factors over the variables this factor shares with them —
+    is computed in
+    :class:`autofit.graphical.mean_field.FactorApproximation`. For most
+    factors that distribution is consumed implicitly: it becomes the
+    prior the search samples from.
+
+    Some hierarchical / population-level analyses want to read those
+    cavity messages directly. A canonical example is a "global"
+    Analysis whose log-likelihood compares model predictions to the
+    per-dataset Gaussian posterior summaries produced by upstream
+    local fits, e.g.::
+
+        log L = -0.5 * sum_i || (pred_i - cavity_mean_i) / cavity_sigma_i ||^2
+
+    To support that, ``EPAnalysisFactor`` attaches the current cavity
+    distribution to its ``Analysis`` immediately before optimisation,
+    via the attribute ``_cavity_mean_field``. The user's
+    ``log_likelihood_function`` can then read each shared variable's
+    cavity message (``.mean`` and ``.scale`` on the
+    ``AbstractMessage`` value) out of the dict.
+
+    The hook is invoked from
+    :func:`autofit.graphical.expectation_propagation.optimiser.factor_step`
+    via duck-typing (``hasattr(factor, "set_cavity_dist")``), so the
+    behaviour of plain ``AnalysisFactor`` is unaffected.
+    """
+
+    def set_cavity_dist(self, cavity_dist):
+        """
+        Store the cavity distribution on the wrapped ``Analysis``.
+
+        Called by :func:`factor_step` once per EP iteration, before this
+        factor's local search runs. The Analysis can read the messages
+        inside ``log_likelihood_function`` by inspecting
+        ``self._cavity_mean_field`` — a ``MeanField`` mapping each
+        shared :class:`Variable` (i.e. ``Prior``) to an
+        ``AbstractMessage`` whose ``.mean`` and ``.scale`` give the
+        cavity Gaussian summary.
+        """
+        self.analysis._cavity_mean_field = cavity_dist

autofit/graphical/expectation_propagation/optimiser.py

@@ -108,6 +108,13 @@ def factor_step(factor_approx, optimiser):
     factor = factor_approx.factor
     factor_logger = logging.getLogger(factor.name)
     factor_logger.debug("Optimising...")
+    # Cavity-message opt-in: factors that implement ``set_cavity_dist``
+    # (e.g. ``EPAnalysisFactor``) receive the current cavity distribution
+    # before optimisation so their Analysis can read per-variable cavity
+    # messages inside ``log_likelihood_function``. Default factors lack
+    # the method, so this is a no-op for them.
+    if hasattr(factor, "set_cavity_dist"):
+        factor.set_cavity_dist(factor_approx.cavity_dist)
     try:
         with LogWarnings(
             logger=factor_logger.debug, action="always"

autofit/graphical/factor_graphs/factor.py

@@ -284,8 +284,6 @@ def _set_jacobians(
         numerical_jacobian=True,
         jacfwd=True,
     ):
-        import jax
-
         self._vjp = vjp
         self._jacfwd = jacfwd
         if vjp or factor_vjp:
@@ -302,8 +300,10 @@ def _set_jacobians(
             elif numerical_jacobian:
                 self._factor_jacobian = self._numerical_factor_jacobian
             elif jacfwd:
+                import jax
                 self._jacobian = jax.jacfwd(self._factor, range(self.n_args))
             else:
+                import jax
                 self._jacobian = jax.jacobian(self._factor, range(self.n_args))
 
     def _factor_value(self, raw_fval) -> FactorValue:

autofit/mapper/prior_model/array.py

@@ -87,15 +87,15 @@ def _instance_for_arguments(
                 pass
 
             if make_array:
-                if isinstance(value, np.ndarray) or isinstance(value, np.float64):
+                if isinstance(value, (np.ndarray, np.float64, float, int)):
                     array = np.zeros(self.shape)
                     make_array = False
                 else:
                     import jax.numpy as jnp
                     array = jnp.zeros(self.shape)
                     make_array = False
 
-            if isinstance(value, np.ndarray) or isinstance(value, np.float64):
+            if isinstance(value, (np.ndarray, np.float64, float, int)):
                 array[index] = value
             else:
                 array = array.at[index].set(value)

autofit/messages/normal.py

@@ -417,7 +417,7 @@ def value_for(self, unit: float) -> float:
         >>> prior = af.GaussianPrior(mean=1.0, sigma=2.0)
         >>> physical_value = prior.value_for(unit=0.5)
         """
-        if isinstance(unit, np.ndarray) or isinstance(unit, np.float64):
+        if isinstance(unit, (np.ndarray, np.float64, float, int, list)):
             from scipy.special import erfinv as scipy_erfinv
             inv = scipy_erfinv(1 - 2.0 * (1.0 - unit))
         else:

autofit/non_linear/analysis/analysis.py

@@ -44,6 +44,30 @@ def __init__(
             use_jax = False
             use_jax_for_visualization = False
 
+        # If the user requested JAX but it isn't installed (e.g. Python <3.11
+        # without the [jax] extra), fall back to numpy with a loud warning
+        # rather than crashing later when the analysis tries to jit-compile.
+        if use_jax:
+            import importlib.util
+            import warnings
+            if importlib.util.find_spec("jax") is None:
+                warnings.warn(
+                    "\n"
+                    "+----------------------------------------------------------------------+\n"
+                    "|  use_jax=True was requested but JAX is not installed.                |\n"
+                    "|                                                                      |\n"
+                    "|  Falling back to numpy. The fit will run, but JAX acceleration       |\n"
+                    "|  (typically 10-100x for large lens models) is unavailable.           |\n"
+                    "|                                                                      |\n"
+                    "|  To enable JAX, install on Python 3.11+ via your library's [jax]     |\n"
+                    "|  extra, e.g.:  pip install autolens[jax]                             |\n"
+                    "+----------------------------------------------------------------------+",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                use_jax = False
+                use_jax_for_visualization = False
+
         if use_jax_for_visualization and not use_jax:
             logger.warning(
                 "use_jax_for_visualization=True requires use_jax=True; "

docs/index.md

@@ -17,7 +17,7 @@ notably [a big data framework for Bayesian hierarchical analysis](https://pyauto
 The following links are useful for new starters:
 
 - [The PyAutoFit readthedocs](https://pyautofit.readthedocs.io/en/latest), which includes an [installation guide](https://pyautofit.readthedocs.io/en/latest/installation/overview.html) and an overview of **PyAutoFit**'s core features.
-- [The introduction Jupyter Notebook on Colab](https://colab.research.google.com/github/PyAutoLabs/autofit_workspace/blob/2026.4.13.6/notebooks/overview/overview_1_the_basics.ipynb), where you can try **PyAutoFit** in a web browser (without installation).
+- [The introduction Jupyter Notebook on Colab](https://colab.research.google.com/github/PyAutoLabs/autofit_workspace/blob/2026.5.1.4/notebooks/overview/overview_1_the_basics.ipynb), where you can try **PyAutoFit** in a web browser (without installation).
 - [The autofit_workspace GitHub repository](https://github.com/Jammy2211/autofit_workspace), which includes example scripts demonstrating **PyAutoFit**'s features.
 - [The standalone HowToFit repository](https://github.com/PyAutoLabs/HowToFit), a series of Jupyter notebook lectures which give new users a step-by-step introduction to **PyAutoFit**.