New ADVI API

pymc_extras/inference/advi/__init__.py

@@ -5,11 +5,38 @@
     AutoMultivariateNormal,
     get_value_shapes_and_dims,
 )
+from pymc_extras.inference.advi.fit import fit_advi
+from pymc_extras.inference.advi.optimizers import (
+    GradientTransformation,
+    adam,
+    apply_updates,
+    chain,
+    clip_by_global_norm,
+    clipped_adam,
+    linear_onecycle_schedule,
+    rmsprop,
+    scale_by_rmsprop,
+    sgd,
+)
+from pymc_extras.inference.advi.training import SVIState, Trainer
 
 __all__ = [
     "AutoDiagonalNormal",
     "AutoGuideModel",
     "AutoLowRankMultivariateNormal",
     "AutoMultivariateNormal",
+    "GradientTransformation",
+    "SVIState",
+    "Trainer",
+    "adam",
+    "apply_updates",
+    "chain",
+    "clip_by_global_norm",
+    "clipped_adam",
+    "fit_advi",
     "get_value_shapes_and_dims",
+    "linear_onecycle_schedule",
+    "rmsprop",
+    "scale_by_rmsprop",
+    "sgd",
 ]

pymc_extras/inference/advi/autoguide.py

@@ -189,7 +189,7 @@ def AutoDiagonalNormal(model: Model, random_seed=None) -> AutoGuideModel:
             )
             Deterministic(
                 rv.name,
-                loc + pt.softplus(scale) * z,
+                loc + pt.exp(scale) * z,
                 dims=value_dims[rv],
             )
 
@@ -252,7 +252,7 @@ def AutoMultivariateNormal(model: Model, random_seed=None) -> AutoGuideModel:
     ordered_rvs = [value_to_rv[name] for name, *_ in point_map_info]
 
     # Initialize the diagonal params at 0.1 (off-diagonal at 0): the full-rank guide starts as
-    # the mean-field guide (diagonal scale softplus(0.1)) and grows correlation structure.
+    # the mean-field guide (diagonal scale exp(0.1)) and grows correlation structure.
     # Matches AutoDiagonalNormal's scale init.
     rows, cols = np.tril_indices(n_dim)
     L_packed_init = np.zeros(rows.size, dtype=loc_init.dtype)
@@ -265,7 +265,7 @@ def AutoMultivariateNormal(model: Model, random_seed=None) -> AutoGuideModel:
         idx = pt.arange(n)
 
         L = pt.zeros((n, n))[pt.tril_indices(n)].set(L_packed)
-        L = L[idx, idx].set(pt.softplus(pt.diagonal(L)))  # positive diagonal
+        L = L[idx, idx].set(pt.exp(pt.diagonal(L)))  # positive diagonal
         # Promise the structure so the MeasurableMatMul logq's solve(L, .) / slogdet(L) lower
         # to the triangular routines instead of a general LU.
         L = assume(L, lower_triangular=True)
@@ -297,7 +297,7 @@ def stochastic_logq(self, path_derivative_gradient: bool = True) -> pt.TensorVar
         u = self.latent
         loc = self["loc"]
         W = self["cov_factor"]  # shape (D, K)
-        d = pt.softplus(self["cov_diag_unconstrained"])  # shape (D,), positive
+        d = pt.exp(self["cov_diag_unconstrained"])  # shape (D,), positive
         n_dim = u.shape[0]
         rank = W.shape[1]
 
@@ -363,15 +363,15 @@ def AutoLowRankMultivariateNormal(
         rank = round(n_dim**0.5)
     rank = max(1, min(rank, n_dim))
 
-    # W starts at 0 (no correlation) and d at softplus(0.1): the guide starts as the mean-field guide.
+    # W starts at 0 (no correlation) and d at exp(0.1): the guide starts as the mean-field guide.
     W_init = np.zeros((n_dim, rank), dtype=loc_init.dtype)
     d_unconstrained_init = np.full(n_dim, 0.1, dtype=loc_init.dtype)
 
     with Model(coords=model.coords, model=None) as guide_model:
         loc = pt.tensor("loc", shape=(None,))
         W = pt.tensor("cov_factor", shape=(None, rank))
         d_unconstrained = pt.tensor("cov_diag_unconstrained", shape=(None,))
-        d = pt.softplus(d_unconstrained)
+        d = pt.exp(d_unconstrained)
         n = loc.shape[0]
 
         eps_k = Normal("eps_k", mu=0.0, sigma=1.0, shape=(rank,))

pymc_extras/inference/advi/compile.py

@@ -1,17 +1,95 @@
 from typing import Protocol
 
 import numpy as np
+import pytensor
 
 from pymc import Model, compile
+from pymc.pytensorf import rewrite_pregrad
+from pytensor import tensor as pt
+from pytensor.compile.sharedvalue import SharedVariable
+from pytensor.graph.replace import graph_replace
 
 from pymc_extras.inference.advi.autoguide import AutoGuideModel
+from pymc_extras.inference.advi.objective import advi_objective, get_logp_logq
+from pymc_extras.inference.advi.optimizers import GradientTransformation
 from pymc_extras.inference.advi.pytensorf import vectorize_random_graph
 
 
+class TrainingFn(Protocol):
+    def __call__(self, *params: np.ndarray) -> tuple[np.ndarray, ...]: ...
+
+
 class SamplingFn(Protocol):
     def __call__(self, *params: np.ndarray) -> tuple[np.ndarray, ...]: ...
 
 
+def compile_svi_step_fn(
+    model: Model,
+    guide: AutoGuideModel,
+    optimizer: GradientTransformation,
+    draws: int = 1,
+    path_derivative_gradient: bool = True,
+    **compile_kwargs,
+) -> tuple[TrainingFn, dict[str, SharedVariable], dict[str, SharedVariable]]:
+    """Compile one full SVI step, with optimizer updates applied in-graph.
+
+    The guide parameters and the optimizer state live in shared variables that the
+    compiled function updates in place. It takes no inputs and returns only the
+    negative ELBO estimate, so no parameters or gradients round-trip through Python
+    during training.
+
+    Together the two returned dicts hold the whole training state: reading their values
+    snapshots a run, writing them resumes one exactly.
+
+    Returns
+    -------
+    step_fn :
+        Compiled function ``step_fn() -> negative_elbo``.
+    shared_params : dict
+        Maps each guide parameter name to the shared variable holding its value.
+    shared_optimizer_state : dict
+        Maps each optimizer state variable name to the shared variable holding its
+        value. Empty for stateless optimizers such as ``sgd``.
+    """
+    if optimizer.pytensor is None:
+        raise ValueError(
+            f"The optimizer {optimizer} does not have a PyTensor implementation "
+            "and cannot be compiled into the step function."
+        )
+
+    logp, logq = get_logp_logq(model, guide, path_derivative_gradient=path_derivative_gradient)
+    scalar_negative_elbo = advi_objective(logp, logq)
+    [negative_elbo_draws] = vectorize_random_graph([scalar_negative_elbo], batch_draws=draws)
+    negative_elbo = negative_elbo_draws.mean(axis=0)
+
+    params_to_shared = {
+        param: pytensor.shared(np.asarray(value), name=param.name)
+        for param, value in guide.params_init_values.items()
+    }
+    [negative_elbo] = graph_replace([negative_elbo], replace=params_to_shared)
+    shared_params = list(params_to_shared.values())
+
+    grads = pt.grad(rewrite_pregrad(negative_elbo), wrt=shared_params)
+
+    new_grads, updates = optimizer.pytensor(grads, shared_params)
+
+    # The optimizer's own state variables are the update keys that are not the guide
+    # parameters themselves.
+    param_ids = {id(param) for param in shared_params}
+    shared_optimizer_state = {var.name: var for var in updates if id(var) not in param_ids}
+
+    for param, grad in zip(shared_params, new_grads):
+        updates[param] = param + grad
+
+    compile_kwargs.setdefault("trust_input", True)
+
+    step_fn = compile(inputs=[], outputs=negative_elbo, updates=updates, **compile_kwargs)
+
+    shared_params_by_name = {param.name: shared for param, shared in params_to_shared.items()}
+
+    return step_fn, shared_params_by_name, shared_optimizer_state
+
+
 def compile_sampling_fn(
     model: Model, guide: AutoGuideModel, draws: int, **compile_kwargs
 ) -> SamplingFn:

pymc_extras/inference/advi/fit.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import numpy as np
+import xarray as xr
+
+from pymc import Model, modelcontext
+from xarray import DataTree
+
+from pymc_extras.inference.advi.optimizers import GradientTransformation
+from pymc_extras.inference.advi.training import Trainer
+
+
+def fit_advi(
+    model: Model | None = None,
+    *,
+    n_steps: int = 10_000,
+    n_particles: int = 1,
+    draws: int = 1_000,
+    optimizer: GradientTransformation | None = None,
+    path_derivative_gradient: bool = True,
+    random_seed=None,
+    backend: str | None = None,
+    compile_kwargs: dict | None = None,
+) -> DataTree:
+    """Fit a model with automatic differentiation variational inference (ADVI).
+
+    Fits a mean-field normal approximation to the model posterior in the unconstrained
+    space, then returns posterior draws from the fitted guide. A one-shot wrapper around
+    :class:`~pymc_extras.inference.advi.training.Trainer` with its default guide; use the
+    trainer directly to keep training, change the optimizer between runs, or sample
+    more than once.
+
+    Parameters
+    ----------
+    model : Model, optional
+        The PyMC model to fit. If None, the model is inferred from context.
+    n_steps : int, optional
+        Number of optimization steps, by default 10_000.
+    n_particles : int, optional
+        Number of guide draws per step used to estimate the ELBO gradient, by default 1.
+    draws : int, optional
+        Number of posterior draws to sample from the fitted guide, by default 1_000.
+    optimizer : GradientTransformation, optional
+        An optax-like optimizer (actual optax optimizers are compatible). By default,
+        :func:`clipped_adam` is used.
+    path_derivative_gradient : bool, optional
+        Whether to use the lower-variance path-derivative ("sticking the landing")
+        gradient estimator, by default True. It is an unbiased variance reduction (it changes
+        only the gradient, not the ELBO); numpyro's ``Trace_ELBO`` does not offer it.
+    random_seed : optional
+        Seed for the guide initialization, the training draws, and the posterior draws.
+    backend : str, optional
+        PyTensor backend to compile the training and sampling functions with
+        (e.g. "numba", "jax", "c"). Mutually exclusive with ``compile_kwargs["mode"]``.
+    compile_kwargs : dict, optional
+        Additional kwargs passed to pytensor compilation.
+
+    Returns
+    -------
+    DataTree
+        Posterior draws from the fitted guide, with the negative loss history in the
+        ``fit`` group (as ``elbo``).
+    """
+    model = modelcontext(model)
+
+    if random_seed is not None:
+        rng = np.random.default_rng(random_seed)
+        init_seed, train_seed, sampling_seed = (int(s) for s in rng.integers(2**30, size=3))
+    else:
+        init_seed = train_seed = sampling_seed = None
+
+    trainer = Trainer(
+        optimizer=optimizer,
+        n_particles=n_particles,
+        path_derivative_gradient=path_derivative_gradient,
+        model=model,
+        backend=backend,
+        compile_kwargs=compile_kwargs,
+        random_seed=init_seed,
+    )
+    state = trainer.fit(n_steps, random_seed=train_seed)
+    idata = trainer.sample_posterior(draws, random_seed=sampling_seed)
+    idata["fit"] = DataTree(dataset=xr.Dataset({"elbo": ("step", -state.loss_history)}))
+    return idata