NumPyro Documentation

NumPyro provides several AutoGuide classes to automatically generate variational distributions (guides) for your models. Use the following guide to select an appropriate one based on your needs:

Basic Mean-Field Guides

Best for starting with variational inference. They automatically handle non-Euclidean latent spaces (e.g., positivity constraints) using bijective transformations.

• AutoNormal: Basic mean-field guide.
• AutoDiagonalNormal: Basic mean-field guide.

Multivariate Normal Guides

These capture correlations in the posterior but may be difficult to fit in high-dimensional settings.

• AutoMultivariateNormal: Full multivariate Normal distribution.
• AutoLowRankMultivariateNormal: Low-rank multivariate Normal distribution.

Normalizing Flow Guides

Provide highly flexible variational distributions parameterized by normalizing flows.

• AutoBNAFNormal: Uses BNAF flows.
• AutoIAFNormal: Uses IAF (Inverse Autoregressive Flow) flows.

Advanced HMC-based Guides

Powerful algorithms that leverage HMC. Good for highly correlated posteriors but can be computationally expensive.

• AutoDAIS: Standard DAIS algorithm.
• AutoSurrogateLikelihoodDAIS: DAIS with support for data subsampling.
• AutoSemiDAIS: Supports data subsampling by using a parametric guide for global latent variables while using DAIS-style approximations for local latent variables.

Other Specialized Guides

• AutoDelta: Used for computing point estimates via MAP (Maximum A Posteriori) estimation.
• AutoLaplaceApproximation: Computes a Laplace approximation.
• AutoGuideList: Used to combine multiple automatic guides.

Unlike Pyro, numpyro.sample requires an explicit random number generator key because JAX does not have a global random state. If you call numpyro.sample outside of an inference context, you have three options:

1. Directly call the distribution and provide a PRNG key: dist.Normal(0, 1).sample(key)

2. Provide the rng_key argument to numpyro.sample: numpyro.sample('x', dist.Normal(0, 1), rng_key=key)

3. Use a seed handler as a context manager or a higher-order function to thread the key automatically.

Using seed as a context manager:

with handlers.seed(rng_seed=0):
    x = numpyro.sample("x", dist.Beta(1, 1))
    y = numpyro.sample("y", dist.Bernoulli(x))

Using seed as a higher-order function:

def fn():
    x = numpyro.sample("x", dist.Beta(1, 1))
    y = numpyro.sample("y", dist.Bernoulli(x))
    return y

print(handlers.seed(fn, rng_seed=0)())

You can install NumPyro using pip or conda.

CPU Installation

To install the latest CPU version of JAX and NumPyro:

pip install numpyro

If you encounter compatibility issues, force a known compatible CPU version:

pip install 'numpyro[cpu]'

GPU Installation

To use NumPyro on a GPU, you must install CUDA first. Use the command corresponding to your CUDA version:

For CUDA 12.x.y:

pip install 'numpyro[cuda12]' -f https://storage.googleapis.com/jax_cuda_releases.html

For CUDA 13.x.y:

pip install 'numpyro[cuda13]' -f https://storage.googleapis.com/jax_cuda_releases.html

Cloud TPU Installation

1. Set up the TPU backend following the Cloud TPU VM JAX Quickstart Guide.
2. Install NumPyro via pip install numpyro.

Other Installation Methods

Conda:

conda install -c conda-forge numpyro

From Source:

git clone https://github.com/pyro-ppl/numpyro.git
cd numpyro
# install jax/jaxlib first for CUDA support
pip install -e '.[dev]'

To perform Bayesian inference using the No-U-Turn Sampler (NUTS), follow these steps:

1. Define your model as a Python callable using numpyro.sample primitives.
2. Initialize the NUTS kernel with your model.
3. Create an MCMC instance specifying the kernel, num_warmup, and num_samples.
4. Execute the inference using mcmc.run(), providing a JAX PRNGKey and any necessary model arguments.

Note: JAX requires explicit PRNG keys for all stochastic operations. Use jax.random.split to manage keys.

Methods available on the MCMC object:

• run(...): Runs warmup, adapts step size/mass matrix, and performs sampling.
• print_summary(): Prints diagnostic information (quantiles, effective sample size, Gelman-Rubin diagnostic).
• get_samples(): Retrieves the posterior distribution samples.
• warmup(...): Performs only the warmup phase.

from numpyro.infer import MCMC, NUTS
from jax import random

# Initialize kernel and MCMC
kernel = NUTS(model)
num_samples = 2000
mcmc = MCMC(kernel, num_warmup=1000, num_samples=num_samples)

# Run inference
rng_key = random.key(0)
rng_key, rng_key_ = random.split(rng_key)
mcmc.run(
    rng_key_, 
    marriage=dset.MarriageScaled.values, 
    divorce=dset.DivorceScaled.values
)

# Retrieve results
mcmc.print_summary()
samples = mcmc.get_samples()