---
id: numpyro
version: "0.21.0"
license: unclear
license_treatment: permissive
maintenance: active
---

# numpyro — Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

License: permissive · Maintenance: active · Downloads: 777.9K/mo

## 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 above — 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

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

- License: not declared (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 777.9K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

probabilistic programming bayesian inference, mcmc hamiltonian monte carlo jax, variational inference automatic differentiation, bayesian statistics sampling, jax-based probabilistic models, nuts sampler inference, hierarchical bayesian modeling, bayesian-inference, jax-accelerated, mcmc-sampling

[View on SkillFed](https://skillfed.io/packages/numpyro) · [View on PyPI](https://pypi.org/project/numpyro/)