Skip to main content

Open-source implementations, interactive demos, and reproducibility materials from my research on Markov chain Monte Carlo methods.

walnutpie

WALNUTS, the Within-Orbit Adaptive Leapfrog No-U-Turn Sampler, adapts the leapfrog step size within each trajectory, so one chain can follow a target whose scale changes from region to region. The reference implementation is walnutpie, a Python package for Markov chain Monte Carlo sampling of differentiable target log densities, maintained at the Flatiron Institute. It takes models from Stan, PyMC, NumPyro, JAX and Numba, or a plain Python function that returns a log density and its gradient.

pip install walnutpie
import walnutpie as wp

def logp(x):
    # your code here
    lp = ...
    gradient = ...
    return lp, gradient

draws = wp.walnuts_pyfunc(logp, num_params=10)

print(wp.ess(draws))

Documentation · Source · PyPI

The core is a header-only C++ library whose only dependency is Eigen. CMake projects can depend on the walnutpie target; otherwise add the include/ directory to your include paths and provide Eigen separately. Code is released under the MIT license, documentation under CC-BY 4.0.

How to cite

N. Bou-Rabee, B. Carpenter, T. S. Kleppe & S. Liu, The Within-Orbit Adaptive Leapfrog No-U-Turn Sampler, Journal of Machine Learning Research, Vol. 27, No. 113, pp. 1-64, 2026.

@article{JMLR:v27:25-1452,
  author  = {Nawaf Bou-Rabee and Bob Carpenter and Tore Selland Kleppe and Sifan Liu},
  title   = {The Within-Orbit Adaptive Leapfrog No-U-Turn Sampler},
  journal = {Journal of Machine Learning Research},
  year    = {2026},
  volume  = {27},
  number  = {113},
  pages   = {1--64},
  url     = {http://jmlr.org/papers/v27/25-1452.html}
}

Demonstrations and reproducibility