Maybe this is obvious, but I’m having trouble figuring out what the advantages and disadvantages of using NumPyro as a backend are. As someone new to the language, how should I decide which to coose? I gather that some MCMC algorithms are faster using NumPyro. Is there anything else to consider?
I think the main benefits of NumPyro are
- Speed
- NumPy api for numerical calculations
- JAX:
- same code for cpu/gpu/tpu
- autobatching
- and much more… (e.g. ode)
Thanks - those are all compelling reasons to use NumPyro. What I’m confused about is if there are cases when one would prefer not to use it (other than it’s Beta status).
Not an expert but the feeling I get from playing with both is that Pyro more focused on SVI methods while numpyro is focusing on HMC and its optimizations.
So if you are doing a lot of BNN Pyro might be the better choice, plus it has been battle tested longer.