MCMC chains with batch dimension when using potential_fn [feature request]

pyro-ppl / pyro

Deep universal probabilistic programming with Python and PyTorch

http://pyro.ai

Apache License 2.0

8.59k stars 988 forks source link

MCMC chains with batch dimension when using potential_fn [feature request] #2539

Open pimdh opened 4 years ago

pimdh commented 4 years ago

Hi,

It seems like having multiple MCMC chains while using potential_fn is currently implemented with just running the chains separately. However, it'd seem more desirable to have an option to compute multiple MCMC chains in parallel using a batch dimension. Tensorflow appears to have this ability (docs) and it'd be great if Pyro could do this too.

Thanks, Pim

fehiepsi commented 4 years ago

Hi @pimdh, currently we don't have a plan to support chain vectorization until PyTorch supports auto-vectorization (which is on the way - from what I heard). You can try NumPyro with chain_method='vectorized', which supports this feature.