exporting the parameters to a model clone in jax/pytorch for speedup
Sounds reasonable to me if you are using some samplers that are slow (like Gamma distribution). If you are using GPU, setting parallel=True will be helpful. If you are using CPU, setting parallel=False will be faster. If you want to distribute the computation across your cores, you can set batch_ndim=0 and use pmap.
def get_pred(sample, rng_key):
predictive = numpyro.infer.Predictive(model, samples, batch_ndim=0)
return predictive(...)
pred = jax.pmap(get_pred)(samples, rng_keys)