feat: analysis-declared gradient_mode (reverse | forward) for gradient searches

autofit/jax/gradient.py

@@ -0,0 +1,118 @@
+"""
+Analysis-declared gradient mode.
+
+An ``af.Analysis`` declares how its likelihood is best differentiated through the class attribute
+``gradient_mode``: ``"reverse"`` (the default, ``jax.value_and_grad``) or ``"forward"``
+(``jax.jacfwd`` over the flat parameter vector). A gradient search may override the declaration
+(e.g. ``af.MultiStartAdam(gradient_mode="reverse")``).
+
+Both modes compute the same gradient; they differ only in cost. Reverse mode costs a small constant
+multiple of one likelihood evaluation whatever the number of parameters, so it is the right default.
+Forward mode costs one tangent per free parameter, but wins when the likelihood itself contains an
+inner forward-mode derivative (reverse mode must then run reverse-over-forward through it) and the
+parameter count is modest -- e.g. the source-plane point-source likelihood, whose lensing Hessian is
+``jax.jacfwd``-built (autolens_profiling #327/#331: 2-4.5x faster, up to 8x faster to compile).
+
+JAX is imported lazily inside the builders so that ``autofit`` imports without JAX.
+"""
+
+from typing import Callable, Optional
+
+GRADIENT_MODES = ("reverse", "forward")
+
+
+def _check_mode(mode: str, source: str) -> str:
+    if mode not in GRADIENT_MODES:
+        raise ValueError(
+            f"Unknown gradient_mode {mode!r} ({source}). Valid modes are "
+            f"{', '.join(repr(m) for m in GRADIENT_MODES)}."
+        )
+    return mode
+
+
+def validate_gradient_mode(mode: Optional[str]) -> Optional[str]:
+    """
+    Return ``mode`` unchanged if it is ``None`` or a valid mode, else raise ``ValueError``.
+
+    Used by constructors (``Fitness``, ``MultiStartGradient``) that accept an optional override,
+    so a mistyped mode fails at construction rather than at the first gradient evaluation.
+    """
+    if mode is None:
+        return None
+    return _check_mode(mode, "override")
+
+
+def resolve_gradient_mode(analysis, override: Optional[str] = None) -> str:
+    """
+    The gradient mode to use for ``analysis``.
+
+    ``override`` (a search- or ``Fitness``-level setting) wins when it is not ``None``; otherwise
+    the analysis's ``gradient_mode`` class attribute is used, defaulting to ``"reverse"`` for
+    objects that do not declare one. An unknown value from either source raises ``ValueError``.
+    """
+    if override is not None:
+        return _check_mode(override, "override")
+
+    mode = getattr(analysis, "gradient_mode", "reverse")
+    return _check_mode(mode, f"declared by {type(analysis).__name__}")
+
+
+def _forward_value_and_grad(func: Callable) -> Callable:
+    import jax
+    import jax.numpy as jnp
+
+    def value_twice(vector):
+        value = func(vector)
+        if jnp.ndim(value) != 0:
+            raise ValueError(
+                "gradient_mode='forward' requires a scalar objective; the function returned "
+                f"an array of shape {jnp.shape(value)}."
+            )
+        return value, value
+
+    jac = jax.jacfwd(value_twice, has_aux=True)
+
+    def value_and_grad(vector):
+        grad, value = jac(vector)
+        return value, grad
+
+    return value_and_grad
+
+
+def value_and_grad_from(func: Callable, mode: str) -> Callable:
+    """
+    A callable returning ``(value, grad)`` of the scalar ``func`` of a flat parameter vector,
+    with the same contract as ``jax.value_and_grad(func)``.
+
+    - ``"reverse"``: ``jax.value_and_grad(func)``;
+    - ``"forward"``: ``jax.jacfwd`` over the vector, with the value carried as auxiliary output so
+      the likelihood is traced once.
+    """
+    _check_mode(mode, "value_and_grad_from")
+
+    if mode == "reverse":
+        import jax
+
+        return jax.value_and_grad(func)
+
+    return _forward_value_and_grad(func)
+
+
+def grad_from(func: Callable, mode: str) -> Callable:
+    """
+    The gradient-only counterpart of ``value_and_grad_from``: ``jax.grad(func)`` in reverse mode,
+    the ``jax.jacfwd`` gradient of the scalar ``func`` in forward mode.
+    """
+    _check_mode(mode, "grad_from")
+
+    if mode == "reverse":
+        import jax
+
+        return jax.grad(func)
+
+    value_and_grad = _forward_value_and_grad(func)
+
+    def grad(vector):
+        return value_and_grad(vector)[1]
+
+    return grad

autofit/non_linear/analysis/analysis.py

@@ -59,6 +59,32 @@ class Analysis(ABC):
     # `ZeroSolver`.
     LATENT_BATCH_MODE = "vmap"
 
+    # How gradient-based searches differentiate this analysis's likelihood
+    # (resolved by `autofit.jax.gradient.resolve_gradient_mode`; used by
+    # `Fitness.grad` and the `MultiStartGradient` searches).
+    #
+    # - "reverse" (default): `jax.value_and_grad`. Its cost is a small constant
+    #   multiple of one likelihood evaluation whatever the number of free
+    #   parameters, so it is right for almost every analysis.
+    #
+    # - "forward": `jax.jacfwd` over the flat parameter vector. Declare it when
+    #   the likelihood itself contains an inner forward-mode derivative (reverse
+    #   mode then runs reverse-over-forward through it) and the model has a
+    #   modest number of free parameters. The source-plane point-source
+    #   likelihood is the reference case: its lensing Hessian is built with
+    #   `jax.jacfwd`, and forward mode measured 2-4.5x faster and up to 8x
+    #   faster to compile through 24 free parameters (autolens_profiling
+    #   #327/#331), so `autolens.AnalysisPoint` declares it.
+    #
+    #   Memory caveat: forward mode carries one tangent per free parameter, so
+    #   under `jax.vmap` (e.g. `MultiStartGradient`'s starts) its memory scales
+    #   with `n_starts * n_params`.
+    #
+    # A search may override the declaration, e.g.
+    # `af.MultiStartAdam(gradient_mode="reverse")`. Both modes compute the same
+    # gradient; only speed and memory differ.
+    gradient_mode = "reverse"
+
     def __init__(
         self,
         use_jax: bool = False,

autofit/non_linear/fitness.py

@@ -15,6 +15,7 @@
 
 
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
+from autofit.jax.gradient import validate_gradient_mode
 from autofit.non_linear.jax_compile import log_on_first_compile
 from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.non_linear.analysis import Analysis
@@ -182,6 +183,7 @@ def __init__(
         iterations_per_quick_update: Optional[int] = None,
         background_quick_update: bool = False,
         live_visual_update: bool = False,
+        gradient_mode: Optional[str] = None,
     ):
         """
         Interfaces with any non-linear search to fit the model to the data and return a log likelihood via
@@ -240,6 +242,10 @@ def __init__(
             chi-squared value that is minimized.
         store_history
             If `True`, the parameters and log likelihood values of every model-fit are stored in lists.
+        gradient_mode
+            How `grad` differentiates the likelihood: `"reverse"` (`jax.grad`) or `"forward"` (`jax.jacfwd`
+            over the flat parameter vector). `None` (default) uses the analysis's declared
+            `Analysis.gradient_mode`. See `autofit.jax.gradient`.
         """
 
         self.analysis = analysis
@@ -264,6 +270,7 @@ def __init__(
 
         self.use_jax_vmap = use_jax_vmap
         self.use_jax_jit = use_jax_jit
+        self.gradient_mode = validate_gradient_mode(gradient_mode)
 
         if getattr(self.analysis, "_use_jax", False):
             from autofit.jax.pytrees import enable_pytrees, register_model
@@ -847,6 +854,8 @@ def __setstate__(self, state):
             or "_apply_assertions_traced" not in self.__dict__
         ):
             self._set_traced_assertions()
+        # `Fitness` objects pickled before the gradient-mode override existed defer to the analysis.
+        self.__dict__.setdefault("gradient_mode", None)
         self._call = self.call
         if getattr(self, "use_jax_vmap", False):
             self._call = self._vmap
@@ -916,22 +925,22 @@ def _jit(self):
     @cached_property
     def _grad(self):
         """
-        Gradient of the JIT-compiled likelihood function.
+        Gradient of the likelihood function (`self.call`) with respect to the flat parameter vector.
 
-        This wraps the JIT-compiled likelihood function (`self._call`) with
-        `jax.grad`, returning a function that computes gradients of the
-        likelihood with respect to its input parameters. Useful for gradient-
-        based optimization and inference methods.
+        The transform is chosen by `autofit.jax.gradient.resolve_gradient_mode`: this `Fitness`'s
+        `gradient_mode` override if set, otherwise the analysis's declared `Analysis.gradient_mode`
+        -- `jax.grad` in reverse mode, `jax.jacfwd` in forward mode. Both give the same gradient.
 
-        Since this is a `cached_property`, the gradient function is compiled
-        and cached on first access, ensuring that expensive setup is done
-        only once.
+        Since this is a `cached_property`, the gradient function is built and cached on first
+        access, ensuring that expensive setup is done only once.
         """
-        import jax
+        from autofit.jax import gradient
+
+        mode = gradient.resolve_gradient_mode(self.analysis, self.gradient_mode)
 
         return log_on_first_compile(
-            jax.grad(self.call),
-            "likelihood function gradient",
+            gradient.grad_from(self.call, mode),
+            f"likelihood function gradient ({mode} mode)",
         )
 
     def grad(self, *args, **kwargs):

autofit/non_linear/search/mle/multi_start_gradient/search.py

@@ -12,6 +12,11 @@
 from autofit.non_linear.search.mle.abstract_mle import AbstractMLE
 from autofit.non_linear.analysis import Analysis
 from autofit.non_linear.fitness import Fitness
+from autofit.jax.gradient import (
+    resolve_gradient_mode,
+    validate_gradient_mode,
+    value_and_grad_from,
+)
 from autofit.non_linear.bijector import AbstractBijector, BijectorNone
 from autofit.non_linear.clipper import (
     AbstractClipper,
@@ -58,6 +63,7 @@ def __init__(
         n_steps: int = 300,
         learning_rate: Optional[float] = None,
         batch_size: Optional[int] = None,
+        gradient_mode: Optional[str] = None,
         max_consecutive_nan: int = 8,
         start_lower_limit: float = 0.15,
         start_upper_limit: float = 0.85,
@@ -142,6 +148,23 @@ def __init__(
             algorithmic knob (how many points it proposes per iteration) that
             autofit merely forwards. Here ``n_starts`` is the algorithm; this
             only decides how many of those starts are evaluated at a time.
+        gradient_mode
+            How the objective is differentiated: ``"reverse"``
+            (``jax.value_and_grad``) or ``"forward"`` (``jax.jacfwd`` over the
+            flat parameter vector). ``None`` (default) uses the analysis's
+            declared ``Analysis.gradient_mode`` -- ``"reverse"`` for every
+            analysis except those that declare otherwise, e.g.
+            ``autolens.AnalysisPoint`` declares ``"forward"`` because its
+            source-plane likelihood carries an inner forward-mode lensing
+            Hessian (autolens_profiling #327/#331). Pass a value to override
+            the declaration for this search. Both modes give the same gradient
+            and so the same fit; only speed, compile time and memory differ.
+
+            Memory caveat, alongside ``batch_size`` above: forward mode carries
+            one tangent per free parameter, so under the vmap over starts its
+            memory scales with ``n_starts * n_params`` (a ``batch_size``-wide
+            chunk's worth when batching). If a forward-mode fit exhausts memory,
+            set ``batch_size`` or override with ``gradient_mode="reverse"``.
         start_lower_limit, start_upper_limit
             The unit-cube bounds broad starts are drawn uniformly from. The
             interior default ``(0.15, 0.85)`` avoids the prior edges where many
@@ -371,6 +394,10 @@ def __init__(
         self.n_starts = n_starts
         self.n_steps = n_steps
         self.batch_size = batch_size
+        # Validated at construction so a mistyped mode fails here rather than
+        # at the first compile; resolved against the analysis at fit time
+        # (``_resolved_gradient_mode``), since no analysis exists yet.
+        self.gradient_mode = validate_gradient_mode(gradient_mode)
         self.max_consecutive_nan = max_consecutive_nan
         self.learning_rate = (
             learning_rate if learning_rate is not None else self._default_learning_rate
@@ -436,6 +463,17 @@ def __init__(
 
         self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search")
 
+    def _resolved_gradient_mode(self, analysis) -> str:
+        """
+        The gradient mode this fit uses: the search's ``gradient_mode`` if set,
+        otherwise the analysis's declaration. ``getattr`` because a search
+        unpickled from before the keyword existed carries no attribute and must
+        defer to the analysis.
+        """
+        return resolve_gradient_mode(
+            analysis, override=getattr(self, "gradient_mode", None)
+        )
+
     def _is_final_boundary(self, converged: bool, total_steps: int) -> bool:
         """
         Whether the chunk boundary just reached ends the search — either the
@@ -971,7 +1009,23 @@ def _fit(
         # compile) while the chunk alone compiles in minutes. The tiling is
         # numerically identical to the vmap; `batch_size` never changes
         # results, it only bounds memory and compile.
-        _value_and_grad = jax.value_and_grad(fitness.call)
+        #
+        # The AD mode (``gradient_mode``) is resolved once here, against the
+        # 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)
+        if not self.silence:
+            source = (
+                "search override"
+                if getattr(self, "gradient_mode", None) is not None
+                else f"declared by {type(analysis).__name__}"
+            )
+            self.logger.info(
+                f"MultiStartGradient gradient mode: {gradient_mode} ({source})."
+            )
+
+        _value_and_grad = value_and_grad_from(fitness.call, gradient_mode)
 
         # Per-start gradient finiteness is reduced *inside* the jitted call, as a
         # third output, rather than by a separate op on its result. The reduction
@@ -1069,7 +1123,9 @@ def _to_physical(vector):
         # objective: ``_broad_starts`` below jits it directly to filter candidate
         # draws, and those draws are, and stay, physical.
         _value_and_grad_stepped = (
-            jax.value_and_grad(lambda phi: fitness.call(_to_physical(phi)))
+            value_and_grad_from(
+                lambda phi: fitness.call(_to_physical(phi)), gradient_mode
+            )
             if (has_scaler or has_bijector)
             else _value_and_grad
         )
@@ -1609,6 +1665,9 @@ def batched_value_and_grad(params):
                 "params": params_physical,
                 "scale": scale,
                 "bijector": self.bijector.kinds if has_bijector else None,
+                # The AD mode this process ran with, resolved against the
+                # analysis (which ``samples_via_internal_from`` cannot see).
+                "gradient_mode": gradient_mode,
                 "opt_state": opt_state,
                 "best_params": best_params,
                 "best_fom": best_fom,
@@ -1983,6 +2042,14 @@ def samples_via_internal_from(
             "n_starts": self.n_starts,
             "n_steps": self.n_steps,
             "batch_size": self.batch_size,
+            # The RESOLVED AD mode (search override, else the analysis's
+            # declaration), read from ``search_internal`` because the analysis
+            # is not available here. Falls back to the search's own setting
+            # (possibly ``None`` = "the analysis's declaration, unrecorded") for
+            # a ``search_internal`` written before the mode was persisted.
+            "gradient_mode": search_internal.get(
+                "gradient_mode", getattr(self, "gradient_mode", None)
+            ),
             "total_steps": total_steps,
             "optax_method": self.optax_method,
             "learning_rate": self.learning_rate,