Adjusted mclmc static

[reubenharry]

Dynamic and static adjusted mclmc

As with HMC, I separate (adjusted) MCLMC into static and dynamic variants, i.e. with and without randomized trajectory length.

I also changed for using sqrt_diag_cov to inverse_mass_matrix (i.e. the square of sqrt_diag_cov) because this is in keeping with the naming scheme for other algorithms. This requires a change in the adaptation scheme and the integrator, but a straightforward one.

This PR is motivated by an upcoming ensemble method.

A few important guidelines and requirements before we can merge your PR:

• If I add a new sampler, there is an issue discussing it already;
• We should be able to understand what the PR does from its title only;
• There is a high-level description of the changes;
• There are links to all the relevant issues, discussions and PRs;
• The branch is rebased on the latest main commit;
• Commit messages follow these guidelines;
• The code respects the current naming conventions;
• Docstrings follow the numpy style guide
• pre-commit is installed and configured on your machine, and you ran it before opening the PR;
• There are tests covering the changes;
• The doc is up-to-date;
• If I add a new sampler* I added/updated related examples

Consider opening a Draft PR if your work is still in progress but you would like some feedback from other contributors.

[junpenglao]

Let me know when it is ready for review

[reubenharry]

@junpenglao Should be ready. Tests seem to be taking a while - not sure why.

[junpenglao]

Seems there are some tests fail

[reubenharry]

Seems there are some tests fail

Hmm, it seems that tests/vi/test_schrodinger_follmer.py::SchrodingerFollmerTest::test_recover_posterior__with_pmap
times out? I'm a bit confused, since this test doesn't seem to involve mclmc or adjusted_mclmc, and I don't believe I made any other changes. Did @AdrienCorenflos write this test, and if so, any ideas?

[junpenglao]

Let me rerun manually

[junpenglao]

hmmm not sure why the test is getting hung up near the end, let me try to delete the cache.

[junpenglao]

hmm same problem that test hang at 95-96%... the reason is not obvious to me. @reubenharry could you close this PR and start afresh? sorry about the trouble

[reubenharry]

Sure, can do!