MCMC in parallel with different datasets

[jeremiecoullon]

Hello!

However I'm wondering if it's possible to run several MCMC in parallel, with each targettng the same posterior with a different dataset. Is this something that is possible with BlackJax, or does the code running on each device have to be the same?

I'm asking as the kernel needs to be built before defining the sampler. So doing a vmap/pmap over different datasets might be possible, but the function being mapped would need to build the kernel..

Thanks!

[junpenglao]

Yes this should be possible with vmapping the inference/sampling routine, something like:

# define a function that runs inference on 1 dataset:
@jax.vmap
def run_inference(dataset):
   ...

the kernel building will encapsulate within the run_inference function

[jeremiecoullon]

Nice! I assumed that the kernel building was not jittable. I'll try this out :)

Thanks very much!