Finally got to opening the issue: Distributed lpdf / grad evaluation · Issue #1425 · pyro-ppl/numpyro · GitHub
I attempted to try your suggested implementation but it was a bit beyond my understanding of the numpyro code base! Thanks again for your help.