Bayesian Inference with NumPyro

NumPyro is a probabilistic programming library built on JAX, giving you Bayesian inference with GPU acceleration and JIT compilation. It's the JAX counterpart to Pyro (which runs on PyTorch).

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, Predictive
import jax
import jax.numpy as jnp

# Define a Bayesian linear regression model
def linear_regression(x, y=None):
    # Priors
    alpha = numpyro.sample('alpha', dist.Normal(0, 10))
    beta = numpyro.sample('beta', dist.Normal(0, 10))
    sigma = numpyro.sample('sigma', dist.HalfNormal(5))

    # Likelihood
    mu = alpha + beta * x
    numpyro.sample('obs', dist.Normal(mu, sigma), obs=y)

# Generate synthetic data
key = jax.random.key(0)
true_alpha, true_beta = 2.0, 3.5
x = jnp.linspace(-5, 5, 100)
y = true_alpha + true_beta * x + 0.5 * jax.random.normal(key, (100,))

# Run MCMC with NUTS (No U-Turn Sampler)
kernel = NUTS(linear_regression)
mcmc = MCMC(kernel, num_warmup=500, num_samples=1000)
mcmc.run(jax.random.key(1), x=x, y=y)

# Get posterior samples
samples = mcmc.get_samples()
print(f"alpha: {samples['alpha'].mean():.2f} ± {samples['alpha'].std():.2f}")
print(f"beta:  {samples['beta'].mean():.2f} ± {samples['beta'].std():.2f}")
# alpha: 2.00 ± 0.05  (true: 2.0)
# beta:  3.50 ± 0.01  (true: 3.5)

# Make predictions
predictive = Predictive(linear_regression, samples)
predictions = predictive(jax.random.key(2), x=jnp.array([0.0, 1.0, 2.0]))
print(f"Predictions: {predictions['obs'].mean(axis=0)}")
# ≈ [2.0, 5.5, 9.0]

Monte Carlo Methods with vmap

import jax
import jax.numpy as jnp

# Monte Carlo estimation of pi using vmap
def estimate_pi(key, num_samples):
    keys = jax.random.split(key, 2)
    x = jax.random.uniform(keys[0], (num_samples,))
    y = jax.random.uniform(keys[1], (num_samples,))
    inside_circle = (x**2 + y**2) <= 1.0
    return 4.0 * jnp.mean(inside_circle)

# Run 100 independent estimates in parallel
keys = jax.random.split(jax.random.key(0), 100)
estimates = jax.vmap(estimate_pi, in_axes=(0, None))(keys, 10000)
print(f"π ≈ {estimates.mean():.4f} ± {estimates.std():.4f}")
# π ≈ 3.1415 ± 0.0162