@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