text: Title: pmap over `num_particles` in SVI
  Body: Hi, 
  
  In `Trace_ELBO`, the `num_particles` argument allows one to effectively introduce a batch size in estimating the ELBO gradient if `num_particles > 1`. By default, it's vectorized over the `num_particles`. Is it possible to also distribute the batch dimension over devices (e.g. when running on multiple GPUs). My particular application is prone to jax OOM errors and would benefit from distribution over `jax.pmap`.