Github1.3万星，迅猛发展的JAX对比TensorFlow、PyTorch

这是一个非常友好的框架，高级 API-Keras 的可用性使得模型层定义、损失函数和模型创建变得非常容易；

TensorFlow2.0 带有 Eager Execution（动态图机制），这使得该库更加用户友好，并且是对以前版本的重大升级；

Keras 这种高级接口有一定的缺点，由于 TensorFlow 抽象了许多底层机制（只是为了方便最终用户），这让研究人员在处理模型方面的自由度更小；

Tensorflow 提供了 TensorBoard，它实际上是 Tensorflow 可视化工具包。它允许研究者可视化损失函数、模型图、模型分析等。

与 TensorFlow 不同，PyTorch 使用动态类型图，这意味着执行图是在运行中创建的。它允许我们随时修改和检查图的内部结构；

除了用户友好的高级 API 之外，PyTorch 还包括精心构建的低级 API，允许对机器学习模型进行越来越多的控制。我们可以在训练期间对模型的前向和后向传递进行检查和修改输出。这被证明对于梯度裁剪和神经风格迁移非常有效；

PyTorch 允许用户扩展代码，可以轻松添加新的损失函数和用户定义的层。PyTorch 的 Autograd 模块实现了深度学习算法中的反向传播求导数，在 Tensor 类上的所有操作， Autograd 都能自动提供微分，简化了手动计算导数的复杂过程；

PyTorch 对数据并行和 GPU 的使用具有广泛的支持；

PyTorch 比 TensorFlow 更 Python 化。PyTorch 非常适合 Python 生态系统，它允许使用 Python 类调试器工具来调试 PyTorch 代码。

正如官方网站所描述的那样，JAX 能够执行 Python+NumPy 程序的可组合转换：向量化、JIT 到 GPU/TPU 等等；

与 PyTorch 相比，JAX 最重要的方面是如何计算梯度。在 Torch 中，图是在前向传递期间创建的，梯度在后向传递期间计算， 另一方面，在 JAX 中，计算表示为函数。在函数上使用 grad() 返回一个梯度函数，该函数直接计算给定输入的函数梯度；

JAX 是一个 autograd 工具，不建议单独使用。有各种基于 JAX 的机器学习库，其中值得注意的是 ObJax、Flax 和 Elegy。由于它们都使用相同的核心并且接口只是 JAX 库的 wrapper，因此可以将它们放在同一个 bracket 下；

Flax 最初是在 PyTorch 生态系统下开发的，更注重使用的灵活性。另一方面，Elegy 受 Keras 启发。ObJAX 主要是为以研究为导向的目的而设计的，它更注重简单性和可理解性。