只知道TF和PyTorch还不够，快来看看怎么从PyTorch转向自动微分神器JAX

选自sjmielke.com机器之心编译机器之心编辑部说到当前的深度学习框架，我们往往绕不开 TensorFlow 和 PyTorch。但除了这两个框架，一些新生力量也不容小觑，其中之一便是 JAX。它具有正向和反向自动微分功能，非常擅长计算高阶导数。这一崭露头角的框架究竟有多好用？怎样用它来展示神经网络内部复杂的梯度更新和反向传播？本文是一个教程贴，教你理解 Jax 的底层逻辑，让你更轻松地从 PyTorch 等进行迁移。Jax 是谷歌开发的一个 Python 库，用于机器学习和数学计算。一经推出，Jax 便将其定义为一个 Python+NumPy 的程序包。它有着可以进行微分、向量化，在 TPU 和 GPU 上采用 JIT 语言等特性。简而言之，这就是 GPU 版本的 numpy，还可以进行自动微分。甚至一些研究者，如 Skye W ………………………………