Nutpie Documentation

nutpie is a high-performance library designed for Bayesian inference, that provides efficient sampling algorithms for probabilistic models. It can sample models that are defined in PyMC or Stan (numpyro and custom hand-coded likelihoods with gradient are coming soon). It provides:

Faster sampling than either the PyMC or Stan default samplers. (An average ~2x speedup on posteriordb compared to Stan)
All the diagnostic information of PyMC and Stan and some more.
GPU support for PyMC models through jax.
A more informative progress bar.
Access to the incomplete trace during sampling.
Experimental normalizing flow adaptation for more efficient sampling of difficult posteriors.

For more details on the algorithms used in nutpie, see the paper Preconditioning Hamiltonian Monte Carlo by minimizing Fisher Divergence. If you use nutpie in your research, please see the citation instructions.

Quickstart: PyMC

Install nutpie with pip, uv, pixi, or conda:

For usage with pymc:

# One of
pip install "nutpie[pymc]"
uv add "nutpie[pymc]"
pixi add nutpie pymc numba
conda install -c conda-forge nutpie pymc numba

And then sample with

import nutpie
import pymc as pm

with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sigma=1)
    obs = pm.Normal("obs", mu=mu, sigma=1, observed=[1, 2, 3])

compiled = nutpie.compile_pymc_model(model)
trace = nutpie.sample(compiled)

Sampler Progress

Total Chains: 6

Active Chains: 0

Finished Chains: 6

Sampling for now

Estimated Time to Completion: now

Draws	Step Size	Gradients/Draw
1400	1.35	1
1400	1.28	3
1400	1.39	3
1400	1.22	3
1400	1.19	1
1400	1.41	1

For more information, see the detailed PyMC usage guide.

Quickstart: Stan

Stan needs access to a compiler toolchain, for which you can find installation instructions here. You can then install nutpie through pip or uv:

# One of
pip install "nutpie[stan]"
uv add "nutpie[stan]"

import nutpie

model = """
data {
    int<lower=0> N;
    vector[N] y;
}
parameters {
    real mu;
}
model {
    mu ~ normal(0, 1);
    y ~ normal(mu, 1);
}
"""

compiled = (
    nutpie
    .compile_stan_model(code=model)
    .with_data(N=3, y=[1, 2, 3])
)
trace = nutpie.sample(compiled)

BridgeStan not found at location specified by $BRIDGESTAN environment variable, downloading version 2.7.0 to /home/runner/.bridgestan/bridgestan-2.7.0
Done!