numpyro

What it is and what it does

NumPyro is a lightweight probabilistic programming library that brings Pyro's API to JAX, enabling fast Bayesian inference through automatic differentiation and JIT compilation. It provides MCMC samplers (including the No-U-Turn Sampler and Hamiltonian Monte Carlo variants), variational inference with flexible guides, a comprehensive distribution library, and effect handlers for building custom inference algorithms.

The library is designed for users building hierarchical Bayesian models, performing posterior inference, or exploring probabilistic programming on modern hardware. It relies on jax, jaxlib, numpy, multipledispatch, and tqdm. The package is actively developed but explicitly warns of potential API changes as the design evolves.

Use it for:

• Run MCMC inference on hierarchical Bayesian models with GPU acceleration via JAX JIT compilation.
• Implement variational inference for models with discrete and continuous latent variables using ADVI.
• Build custom inference algorithms by composing effect handlers and Pyro primitives.
• Perform Bayesian data analysis with standard distributions and constraints similar to PyTorch's API.
• Accelerate Hamiltonian Monte Carlo by compiling the entire verlet integrator and tree-building stage.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

NumPyro is a probabilistic programming library that uses JAX for automatic differentiation and JIT compilation, enabling Bayesian inference with MCMC and variational inference algorithms on CPU, GPU, and TPU.

Yes, if you need probabilistic programming with GPU/TPU acceleration and are comfortable with an actively-developed library that may change its API. The low install friction, active maintenance, permissive license, and strong maintenance signal (recent release, 2736 stars) make it a solid choice for Bayesian inference workflows. Not recommended if you require API stability or are new to probabilistic programming.

Install

numpyro on PyPI

pip install numpyro

uv add numpyro

poetry add numpyro

Installing numpyro

Before you install

Installation is straightforward with low friction; the package is actively maintained with a recent release and carries 5 runtime dependencies including jax and jaxlib. The codebase is under active development, so API stability is not guaranteed.

License in practice

Licensed under Apache 2.0 (permissive), allowing free use, modification, and distribution with minimal restrictions.

Quickstart

pip install numpyro

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
from jax import random

def model(data):
    mu = numpyro.sample('mu', dist.Normal(0, 1))
    numpyro.sample('obs', dist.Normal(mu, 1), obs=data)

kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=1000)
mcmc.run(random.key(0), data)

Requires JAX and jaxlib, which have their own system-level dependencies (CUDA/cuDNN for GPU support); Python 3.11 or later.

Verify before relying

• Whether the API brittleness mentioned in the description excerpt affects common use cases or only edge cases.
• Performance characteristics compared to other probabilistic programming frameworks for typical model sizes.

Package facts

Evidence: numpyro-0.21.0-py3-none-any.whl

Keywords: probabilistic, machine learning, bayesian, statistics

Tags

More Artificial Intelligence packages