Adaptation support / example for multiple chains

Dear blackjax team,

Thank you for your hard work, this project is really cool!

I’ve been playing around with the PyMC support, and I think I’ve managed to update the example notebook to use PyMC v4 and aesara. This wasn’t very hard to do thankfully, and I’ll be happy to contribute it as a new example / update the old one if that’s of interest.

However, my ultimate aim would be to include blackjax as part of the recent MCMC speed comparisons (see code and blog post). To get the best out of blackjax, I’d like to run four chains in parallel like with the other methods. That’s where I’ve run into a bit of trouble! There’s an example of sampling multiple chains in the Intro notebook here but I don’t see anything about also doing separate adaptation…! I’m actually not sure what the best way to go is there – should chains be adapted separately or one step size and mass matrix should be estimated and shared in a warmup (probably)?

In any case, I feel like an example of how best to do adaption with multiple chains would be helpful to add. Please let me know if I missed an existing example somewhere, or if there is already a recommended approach.

Thanks for your help, Martin

[rlouf]

You’re welcome! You’ll hopefully find the user interface more pleasant as well 😃

[rlouf]

Hey @martiningram lots of work since I last replied on #153 and #159. We’ve identified one major difference with Numpyro which likely explains the perf difference, will keep you posted!