Posterior predictive on new data with MCMC in a model with local and global parameters

I have a latent factor model with global parameters and (local) latent parameters, see twitter thread here

I want to do inference with MCMC, marginalizing out the latent variable to get the correct parameter estimates. Then, to do posterior predictions on new data, I need to re-run MCMC multiple times on the test data, each time using a different sample from the posterior for the global parameters, but using MCMC to again marginalize out the local latent variables. I’ve spent a lot of time trying to do this with (num)pyro’s effect handlers, but I haven’t found a way to do it the ‘pyro’ way. Can someone give me a hand?

Here’s a self contained example: