@hellocyber Recently, we build a numpy (based on JAX) backend for Pyro. I just add a test for the above model at Test some edge examples from Pyro by fehiepsi 路 Pull Request #131 路 pyro-ppl/numpyro 路 GitHub. It is very fast (100x faster than Pyro for this model) so I guess that you want to try it out. 
2 Likes