Utilities for MCMC about numpyro

[neerajprad]

@fehiepsi - I'll try to hook up a basic HMC implementation using the utilities that you already have put in place, just to explore how the API looks like, e.g. use of vmap for doing batching.

[fehiepsi]

@neerajprad Sure, please do it. While implementing the above functions, the trickiest part to me is to handle state such that jit version will work correctly. It is complicated when there are many if/else statements in our functions. I have been getting stuck with warmup_adapter for a while due to many conditions need to handle.

[neerajprad]

That is indeed quite tricky and I think NUTS will be even trickier, and will need a lot of thinking. :) I'll start simple and take a look at how far I get. I don't think we need to aim for a 100% jitted algorithm though, even just jiting the model PE and backwards will take us 80% there.