How to move MCMC run on GPU to CPU

numpyro
4

What do you mean with " there are some missing jax.device_get logic in arviz "?

Should this be in Arviz? Or should the numpyro.infer.mcmc.MCMC object have a device_put method?