MCMC Support#

While aimz is primarily designed around variational inference and predictive sampling, it also provides support for MCMC methods via the NumPyro backend, using the same aimz interface (e.g., fit_on_batch() and predict_on_batch()). This enables users to apply MCMC to more complex models where variational inference may be less effective and dataset sizes are relatively small.

import logging

import jax.numpy as jnp
import numpyro.distributions as dist
from jax import random
from jax.typing import ArrayLike
from numpyro import plate, sample
from numpyro.infer import MCMC, NUTS

from aimz import ImpactModel

logging.basicConfig(level=logging.INFO, force=True)

Model and Data#

We set up a linear regression model and create synthetic data for both features and targets as an example.

def model(X: ArrayLike, y: ArrayLike | None = None) -> None:
    """Linear regression model."""
    w = sample("w", dist.Normal().expand((X.shape[1],)))
    b = sample("b", dist.Normal())
    mu = jnp.dot(X, w) + b
    sigma = sample("sigma", dist.Exponential())
    with plate("data", size=X.shape[0]):
        sample("y", dist.Normal(mu, sigma), obs=y)


rng_key = random.key(42)
rng_key, rng_key_w, rng_key_b, rng_key_x, rng_key_e = random.split(rng_key, 5)
w = random.normal(rng_key_w, (10,))
b = random.normal(rng_key_b)
X = random.normal(rng_key_x, (1000, 10))
e = random.normal(rng_key_e, (1000,))
y = jnp.dot(X, w) + b + e

MCMC Support#

Model and Data#

MCMC Sampling and Prediction#

Using External MCMC Samples#