Implement ChEES-HMC

I have a working draft of ChEES-HMC, except its not JAX-ready yet. The “Select jittered trajectory length…” line of Algorithm 1 requires a float from a predetermined Halton sequence of length wamup+sampling number of iterations. When returning the kernel after warm-up, it needs to keep yielding from the sequence at every iteration (but JAX functions need to be pure).

For now my implementation yields the next number using a scipy:

halton_sequence_gn = scipy.stats.qmc.Halton(d=1)
#then inside the kernel...
currrent_halton_number = jnp.array(halton_sequence_gn.random(n=1).squeeze())

Any help is appreciated. Will keep updated on progress.

[albcab]

Have it almost done, just some minor bugs. I’ve been pretty busy these past few months, but I have some time today so I think I can get it done 👍

[rlouf]

That works too! I know the PRNG algorithm JAX uses pretty well and there’s no easy way to engineer an int sequence from it; plus since the keys are provided by the user you would get no guarantee on the ordering.

You could also do like we do for window adaptation, which is scaning over a pre-generated sequence (schedule for the window adaptation). That’s much cleaner than carrying an integer. That would be acceptable at the lower level (in mcmc.chees).