简单地说，JAX、Trax和TensorRT之间有什么区别？

Question

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

然后我听说：

日本航空公司https://github.com/google/jax

• Trax https://github.com/google/trax

两者似乎都加速了DL。但我现在很难理解他们。有人能用简单的术语来解释吗？

Answer 1

Trax是谷歌创建的一个深度学习框架，被谷歌大脑团队广泛使用。当涉及到实现最先进的深度学习模型的现成状态时，它作为TensorFlow和PyTorch的一种替代，例如，在自然语言处理领域，原则上是Transformers，Bert等。

Trax建立在TensorFlow和JAX的基础上。JAX是Numpy的增强和优化版本。JAX和NumPy的重要区别在于前者使用了一个名为XLA (高级线性代数)的库，它允许在GPU和TPU上运行NumPy代码，而不是像在普通NumPy中那样在CPU上运行代码，从而加快了计算速度。