Blackjax MCLMC support

[reubenharry]

MCLMC is an algorithm intended as a statistically more efficient alternative to NUTS HMC. I'm interested in adding the option to PyMC to use it via the Blackjax implementation, in the spirit of pm.sample(nuts_sampler="blackjax"), so maybe something like pm.sample(mclmc_sampler="blackjax").

It comes with either an unadjusted version (no MH, controlled asymptotic bias, faster) or an MH adjusted version, both of which are faster than NUTS for a range of benchmarks of interest.

Basically my question is: would PyMC be open to a PR along these lines? For public facing visibility, we'd like to put the algorithm into PyMC rather than use PyMC indirectly (i.e. extract a density from a probabilistic program in PyMC), and I thought this might be the best way.

[ricardoV94]

Is bayeux a better tool to allow combining pymc with blackjax samplers? CC @ColCarroll

sample is already a very overloaded function, so I wouldn't want to add a new mclmc_sampler kwarg to it.

[ricardoV94]

Needs someone to have a go at it. I don't have the complete mental map to anticipate edge cases. First guess is that it may be a more scalable solution.

An external sampler class and a monolithic external sampler subclass that tell pymc stuff will be sampled externally and must sample the whole model. The external nuts samplers should be implemented like this then as well

[twiecki]

@reubenharry This is probably a bit more than what you bargained for I assume?

[twiecki]

Claude Code at a go at the interface for nutpie in #7719

[reubenharry]

Thanks for this discussion! Not necessarily more than I bargained for, but might take me some time to understand how to approach (perhaps with Claude's help, as you suggest).

[ricardoV94]

We're happy to help, this is something that would be nice to have so its easier to plug alternative samplers in the future

[reubenharry]

Just to give a sense of the performance benefits, here's our paper, with a comparison between NUTS and Metropolis Adjusted Microcanonical Sampling (MAMS) on some standard test problems, in table 1: https://arxiv.org/pdf/2503.01707.

MAMS is nowhere near as battle-tested as NUTS, but it seems to give a 2x or greater speedup in most cases we tried, is easier to parallelize, and we're interested in getting input from pymc users with difficult problems (hence this discussion).

[ricardoV94]

Ricardo wanted

Ricardo wants* :D It's just too good not to

[reubenharry]

Ah yes, I had forgotten that. But yes, it would be nice to try obtaining the nutpie mass matrix or normalizing flow for a given problem, and then subsequently run MAMS instead of NUTS. Our trajectory length tuning is described in section 5.3 of the paper linked above. (It could actually also be used for HMC, but we haven't tried that).

[aseyboldt]

nutpie.sample(model, store_mass_matrix=True).sample_stats.mass_matrix_inv :-)

[ColCarroll]

[aseyboldt]

Sorry Colin 😅