Software
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
- WALNUTS interactive demo
Watch adaptive step-size sampling explore a ring, funnel, banana or Gaussian, live in the browser. Built by Jeremy Magland. - NUTS vs WALNUTS animation
Split-screen comparison of the two samplers on Neal’s funnel, with code to reproduce it. - github.com/nawafbourabee
Additional code and reproducibility materials.