Fix xla_pmap_p import for JAX versions that removed pmap#2173

juanitorduz merged 2 commits into

pyro-ppl:masterpyro-ppl/numpyro:masterfrom

saitcakmak:fix-xla-pmap-p-importsaitcakmak/numpyro:fix-xla-pmap-p-importCopy head branch name to clipboard