简单来说，JAX、Trax 和 TensorRT 有什么区别？

Question

我一直在使用 TensorRT 和 TensorFlow-TRT 来加速我的 DL 算法的推理。

然后我听说了：

• JAX https://github.com/google/jax
• 特拉克斯https://github.com/google/trax

两者似乎都加速了深度学习。但我很难理解他们。谁能简单的解释一下？

Answer 1

Trax 是由谷歌创建并被谷歌大脑团队广泛使用的深度学习框架。当涉及到实施现成的最先进的深度学习模型（例如 Transformers、Bert 等）时，它是 TensorFlow 和 PyTorch 的替代方案，原则上与自然语言处理领域相关。

Trax 建立在 TensorFlow 和 JAX 之上。 JAX 是 Numpy 的增强和优化版本。 JAX 和 NumPy 的重要区别在于前者使用称为 XLA（高级线性代数）的库，它允许在 GPU 和 TPU 上运行 NumPy 代码，而不是像这样在 CPU 上运行发生在平原NumPy，从而加快计算速度。