Hi everyone,
First of all, I would like to point out that I’m fairly new to jax and numpyro, so I’m sorry if my questions are be a bit naive.
I would like to do inference with MCMC with a Bayesian model f_\theta parametrized by \theta^{(s)} \sim p(\theta|\left\{x^\mathrm{train},y^\mathrm{train}\right\}) trained such that y \approx f_\theta(x). I have an observation y^\mathrm{obs} and I want to find associated x^\mathrm{obs}\sim p(x^\mathrm{obs}|y^\mathrm{obs})\propto p(y^\mathrm{obs}|x^\mathrm{obs})p(x^\mathrm{obs}) using MCMC with f_\theta used in the computation of the likelihood p(y^\mathrm{obs}|x^\mathrm{obs}). The thing is the model f_\theta is itself Bayesian, and I have multiple samples for \theta. According to Reiser et al. 2025 I can marginalize over \theta and compute the expected posterior (E-post) distribution as
which can be approximated using S independent MCMC, one for each model parameter sample \theta^{(s)}. As they are independent, I would like to run them in parallel. What is the correct way to do it? Is it possible to use vmap/pmap over \theta^{(s)}? Or must I run them sequentially with a for loop over \theta^{(s)}?
Thank you!