We are excited to announce the first release of NumPyro. NumPyro provides a NumPy backed for Pyro using JAX for automatic differentiation and JIT compilation to GPU and CPU. It currently supports the same modeling API as Pyro with first class support for HMC using the No-U-Turn Sampler (NUTS). This is an alpha release, and there is more to come. Please give it a try and let us know what you think.
~ Pyro Dev Team.