WALNUTS: A New Windowed Adaptive Leapfrog NUTS Kernel

[DJLacombeTTU]

Hi everyone,

Following up on the Slingshot sampler I shared a few days ago, I’m excited to introduce another custom transition kernel I’ve been developing for BlackJax: WALNUTS (Windowed Adaptive Leapfrog with NUTS).

This sampler is a pure JAX adaptation of the WALNUTS algorithm recently proposed by Bou-Rabee et al. (2025) in their paper, "The Within-Orbit Adaptive Leapfrog No-U-Turn Sampler". They established the theoretical foundation for adapting the leapfrog step size at fixed intervals of simulated time as the orbit evolves, while preserving reversibility through a carefully designed extended state space and measure-preserving involution.

As an applied econometrician, I frequently run into deep, multiscale hierarchical structures—specifically spatial autoregressive models and flexible income estimations where the group-level variances dictate lower-level parameters. Standard NUTS often struggles with these highly correlated geometries, either taking premature U-turns or requiring massive, computationally expensive dense mass matrices to avoid divergence.

WALNUTS is designed to tackle this by integrating dynamic micro-halving directly into the trajectory building, autonomously navigating steep gradients without permanently shrinking the macro step size, as detailed in the original paper.

Here are the technical highlights of my BlackJax implementation:

• Pure JAX & Stateless Architecture: The kernel is fully compliant with BlackJax's functional API (blackjax/mcmc/walnuts.py), making it perfectly safe for jax.jit and multi-GPU vectorization via jax.pmap/jax.vmap.
• Decoupled Adaptation: I built a custom, expanding-window adaptation schedule (blackjax/adaptation/dense_window_adaptation.py) that handles Dual Averaging and Dense Welford covariance tracking. It strictly bounds the leapfrog integrator to prevent resonance traps in standardized spaces.
• Empirical Validation: I stress-tested the kernel natively in JAX against a 10-dimensional Neal’s Funnel. By deploying a fully isotropic non-centered parameterization, the engine effortlessly mapped the hypersphere. Running 4 parallel chains (8,000 pooled draws) yielded flawless stationary convergence (all r_hat = 1.0), massive Effective Sample Sizes, and nailed the true variance target without a single divergent transition.
• PyMC Compatibility: The architecture bridges seamlessly with PyMC. I’ve successfully mapped PyTensor graphs to XLA and recovered exact parameter distributions with high efficiency.

I’d love to get the maintainers' eyes on the internal bounds I set for the Dual Averaging and the XLA typings in the Welford state reset.

Looking forward to hearing your thoughts and seeing how this might fit into the broader BlackJax ecosystem!

Reference:
Bou-Rabee, N., Carpenter, B., Kleppe, T. S., & Liu, S. (2025). The Within-Orbit Adaptive Leapfrog No-U-Turn Sampler. arXiv preprint arXiv:2506.18746v1.

[DJLacombeTTU]

I forgot to add the repo to the last message: https://github.com/DJLacombeTTU/WALNUTS

[DJLacombeTTU]

I'm working on getting this to be 100% compatible with blackjax and hopefully will be smoother than the previous task.