ArchestackBest of Open-Source AI: List on GitHub ↗ (opens in a new tab)

#131131st of 214 in ML frameworks and libraries

numpyro

Probabilistic programming library on JAX with NUTS/HMC and variational inference

Documentation ↗ (opens in a new tab)Website ↗ (opens in a new tab)Repository on GitHub ↗ (opens in a new tab)
Stars
2.8k
License
Apache-2.0
Last commit
Nov 2026
Last release
Sep 2026
Language
Python

Missed the latest GitHub check; these facts may be older.

Overview

NumPyro is a Python probabilistic programming library that provides a NumPy/JAX backend for Pyro's model API (sample, param, plate, effect handlers). It JIT-compiles inference, including the full NUTS tree-building step, through XLA to CPU, GPU or TPU. Inference covers MCMC variants (NUTS, HMC, MixedHMC, HMCECS, BarkerMH, SA) and stochastic variational inference with autoguides.

Who it is for: Statisticians and ML researchers doing Bayesian inference in Python on JAX

Strengths

  • JIT-compiles the whole NUTS tree build with JAX, removing Python overhead
  • Broad MCMC set: NUTS, HMC, MixedHMC, HMCECS, BarkerMH, HMCGibbs, SA
  • Pyro-like API and torch.distributions-style distributions; TFP distributions usable in models
  • Automatically sums out finite discrete latents so NUTS/HMC can run on the rest

Weaknesses

  • README warns of brittleness, bugs and API changes under active development
  • No global param store or random state; models must be written functionally
  • Requires explicit PRNG key handling, unfamiliar to PyTorch/Pyro users
  • Inference gets harder as latent dimension grows; may need reparameterization

What it needs

  • GPU optional
  • Needs JAX
  • Models: NUTS, HMC, MixedHMC, HMCECS, BarkerMH

Also in ML frameworks and libraries

See all 214
Also in ML frameworks and libraries
RankProjectScore
1transformersPython library defining open models for text, vision, audio and multimodal tasks168.9k stars+1,573 (+0.9%) in 30 daysApache-2.094 out of 100
2ultralyticsPython library and CLI for training and running YOLO vision models63.3k stars+1,004 (+1.6%) in 30 daysAGPL-3.094 out of 100
3pytorch-image-modelsPyTorch library of image models, pretrained weights and training scripts37.3k stars+107 (+0.3%) in 30 daysApache-2.094 out of 100
#4pytorchPython tensor library with GPU support and tape-based autograd105.5k stars+1,036 (+1.0%) in 30 dayscustom license93 out of 100
#5supervisionPython toolkit for loading, annotating and counting computer vision detections Live demo ↗ (opens in a new tab)51.6k stars+521 (+1.0%) in 30 daysMIT92 out of 100