Adding new samplers to NumPyro

[reubenharry]

Feature Summary

Recently, a couple of samplers have been proposed by the group I work with (https://arxiv.org/abs/2503.01707, https://arxiv.org/abs/2212.08549) as alternatives to NUTS HMC (see issue #1662
). Since they appear to be quite a bit faster than NUTS (at least on benchmark problems I've tried), and relatively simple, I'm interested in adding them to NumPyro, but wanted to get some advice.

Currently, implementations exist in Blackjax. In an ideal world, I'd make a new class like class AdjustedMicrocanonical(numpyro.infer.mcmc.MCMCKernel) which basically just wraps Blackjax.

In addition, my eventual goal would be to add not just the kernel, but also the tuning scheme (which is key to good performance). I'm curious if there's a straightforward way to do that.

Motivation

While it's easy to write a model in NumPyro and extract the density, then use Blackjax for inference, we want to give users more direct access (basically for the purpose of increasing discoverability).

[martinjankowiak]

one issue with relying on a blackjax implementation is that the numpyro implementation will tend to break every so often as blackjax is refactored. that has certainly been our experience with wrapping jaxns => https://num.pyro.ai/en/latest/contrib.html#nested-sampling

[reubenharry]

Noted. I maintain the blackjax samplers, and would plan to add a stable one (we haven't changed the unadjusted MCLMC sampler for a year or so), so that might offset the problem somewhat.

What I'm more unsure of is what it would look like to add a hyperparameter adaptation algorithm to numpyro - is this feasible?

Edit: on further inspection, it seems like it is, so I'll try to make something like the existing HMC (but for MCLMC), which depends on blackjax, and then we can discuss this dependency down the line.

[fehiepsi]

Just curious, is it easier to use blackjax for inference with those samplers? If so, I think you can add a section to https://num.pyro.ai/en/stable/tutorials/other_samplers.html to illustrate the usage.

[reubenharry]

Yes, it's easier, but my motivation for adding samplers directly to Numpyro is to make them more accessible (and in particular, discoverable) for users of Numpyro.

[martinjankowiak]

also, depending on how it's done, you could get things like discrete enumeration for "free"

[reubenharry]

Yeah, I did this in a pretty naive way in #2039, following the MCMCKernel class, because the HMC code was pretty daunting, but could be improved in future.