pytorch - NumPyro 与 Pyro : Why is former 100x faster and when should I use the latter?

Question

来自 Pytorch-Pyro 的 website :

We’re excited to announce the release of NumPyro, a NumPy-backed Pyro using JAX for automatic differentiation and JIT compilation, with over 100x speedup for HMC and NUTS!

我的问题:

NumPyro(超过 Pyro)的性能增益(即 sometimes 340 倍或 2 倍)究竟来自哪里？

更重要的是，为什么(而不是，在哪里)我会继续使用 Pyro？

额外的:

与 Tensorflow Probability 相比，我应该如何查看 NumPyro 的性能和特性，以决定在何处使用哪个？

Answer 1

这是个好问题。我刚刚在 Pyro's dedicated forum 中问了同样的问题.这是他们的一位核心开发人员的回答:“Pyro 中有很多很酷的东西在 NumPyro 中没有出现，例如，请参阅 Pyro 文档中的贡献代码部分。对我来说，在开发过程中，调试 PyTorch 代码要容易得多比 Jax 代码(尽管 Jax 团队在最近的版本中付出了很多努力来帮助调试)。因此，为了实现新的推理算法，我在 Pyro 中工作更容易。”