feat: Add Microcanonical Langevin Monte Carlo (MCLMC) kernel

[juanitorduz]

Trying to support #2039

[juanitorduz]

Hey @reubenharry I tried using your branch, but I made a git mess 🙈 so I decided to open another one to get feedback (I will add you as a coauthor).

Could you please check this one out and see if the implementation (and the tests) are as expected?

[fehiepsi]

I worry that we might introduce technical debt by depending on other libraries. Could we turn this into an example/tutorial instead? It's not clear to me the benefit of using numpyro here.

[fehiepsi]

Maybe adding a section for https://num.pyro.ai/en/stable/tutorials/other_samplers.html would be better?

edit: #2035 has good discussion about the above points

[juanitorduz]

I worry that we might introduce technical debt by depending on other libraries. Could we turn this into an example/tutorial instead? It's not clear to me the benefit of using numpyro here.

Sure! This was a first attempt at trying to see how it would fit. I agree this "optional" dependencies can be hard to maintain.

So what about adding another section with this code in https://num.pyro.ai/en/stable/tutorials/other_samplers.html ?

Or do we want an additional notebook with just a brief explanation of MCLMC?

[martinjankowiak]

on the other hand the optimal outcome would be that its very easy for numpyro users to use the sampler. if it's hidden in some tutorial...

[juanitorduz]

on the other hand the optimal outcome would be that its very easy for numpyro users to use the sampler. if it's hidden in some tutorial...

True 😄 . I do not have any strong opinion. I just wanted to bring this one alive. That being said, I have seen blackjax making breaking changes and this would be a pain to maintain indeed (in the notebook we would put a disclaimer about these potential changes)

[juanitorduz]

ok! In af9f001 I added blackjax to the tests-deps just to ensure the tests are green.

After that I can try to port it into the notebook (and remove any blackjax dependency) for now and keep the discussion open in #2035 . We can always come back to the commit history and put it back in NumPyro :)

[martinjankowiak]

how entangled is the blackjax implementation with the rest of blackjax? can the sampler be ripped out with minimal changes?

[reubenharry]

Awesome, thanks for doing this! It looks good to me, although for peace of mind, we should take a non-trivial example and check that the efficiency is good (e.g. Stochastic Volatility or something). Perhaps I can do that by running the Numpyro version in my https://github.com/reubenharry/sampler-benchmarks repo and checking the results, when I have more time.

@martinjankowiak Re entanglement, I have previously ported the implementation outside Blackjax before (even out of Jax), and it isn't insanely hard to do so (blackjax itself isn't really a very complex codebase). More a question of time. I agree that the main draw of adding the code to Numpyro is discoverability for users.

(I have been a bit busy this month, but will try to be more responsive going forward)

[reubenharry]

This implementation works for me. I'd say that if it was totally necessary, the relevant parts of blackjax could be pulled out and the implementation could be self-contained. But I think it would be more natural to keep it in - I think changes that break backwards compatibility of this sampler are unlikely to be frequent.

[juanitorduz]

@fehiepsi @martinjankowiak @reubenharry

Suggestion: Keep this code but put it into contrib.infer and make blackjax as an optional dependency as in https://github.com/pyro-ppl/numpyro/blob/master/numpyro/contrib/nested_sampling.py :

numpyro/numpyro/contrib/nested_sampling.py

Lines 10 to 29 in d5598e7

	try:
	import jaxns # noqa: F401
	from jaxns import (
	Model,
	Prior,
	TerminationCondition,
	plot_cornerplot,
	plot_diagnostics,
	resample,
	summary,
	)
	from jaxns.public import DefaultNestedSampler
	from jaxns.utils import NestedSamplerResults
	except ImportError as e:
	raise ImportError(
	f"{e} \n "
	f"To use this module, please install `jaxns>2.5` package. It can be"
	" installed with `pip install jaxns` with python>=3.8"
	) from e

[fehiepsi]

Hi @juanitorduz, I planned to remove jaxns wrapper at some point to avoid maintenance overhead. So I hope we can avoid adding wrappers for other libraries.

[martinjankowiak]

it seems to me like the best option is to pull out the implementation from blackjax. this is something claude could presumably do in < 10 minutes. the only downside is that you don't automatically inherit any improvements made to the blackjax implementation, but that's not a big deal and can always be addressed with future PRs if there are significant improvements in the blackjax implementation.

[juanitorduz]

ok! Fair enough! Then I will proceed to pull out the implementation from blackjax :) 👍

[juanitorduz]

@reubenharry I have tried to remove the blackjax dependency with Claude's help! I ensured the results match the code from #2039 but it would be great if you could give a proper review

@fehiepsi @martinjankowiak Open for feedback :)

[martinjankowiak]

presumably many of the private methods in this file have corresponding tests in blackjax that you can pull into this PR?

[juanitorduz]

Added in e2d754c

[juanitorduz]

@reubenharry I have tried to remove the blackjax dependency with Claude's help! I ensured the results match the code from #2039 but it would be great if you could give a proper review

@fehiepsi @martinjankowiak Open for feedback :)

hey folks! any thoughts on this one? This raw implementation is self-contained and about 859 lines of code (I am fine with it)

I wanna also want to bring up the fact that blackjax is quite active now https://github.com/blackjax-devs/blackjax/releases, so my gut tells me it's less risky to have it as an optional dependency (much safer than jaxns).

[reubenharry]

Sorry for the slow reply. I had one minor comment (see above), but since this follows the BlackJax quite closely, the rest looks fine. I think either the self contained version, or using Blackjax directly would be fine (but if the latter, I'd suggest that NumPyro makes a larger scale decision to be open to using Blackjax across the board).