pytorch - NumPyro vs Pyro：为什么前者快 100 倍，我什么时候应该使用后者？

Question

来自 Pytorch-Pyro 的网站：

我们很高兴地宣布发布 NumPyro，这是一个由 NumPy 支持的 Pyro，使用 JAX 进行自动微分和 JIT 编译，HMC 和 NUTS 的加速超过 100 倍！

我的问题：

1. NumPyro（超过 Pyro）的性能增益（有时是 340 倍或 2 倍）究竟来自哪里？
2. 更重要的是，为什么（更确切地说，在哪里）我会继续使用 Pyro？

额外的：

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

Answer 1

这是个好问题。我刚刚在Pyro 的专门论坛上问了同样的问题。这是他们的一位核心开发人员的回答：“Pyro 中有很多很酷的东西没有出现在 NumPyro 中，例如，请参阅 Pyro 文档中的贡献代码部分。对我来说，在开发时，调试 PyTorch 代码要容易得多“