出处丨AI前线

什么？TensorFlow 有了替代品？什么？竟然还是谷歌自己做出来的？先别慌，从各种意义上来说，这个所谓的“替代品”其实是 TensorFlow 的一个简化库，名为 JAX，结合 Autograd 和 XLA，可以支持部分 TensorFlow 的功能，但是比 TensorFlow 更加简洁易用。虽然还不至于替代 TensorFlow，但已经有 Reddit 网友对 JAX 寄予厚望，并表示“早就期待能有一个可以直接调用 Numpy API 接口的库了！”，“希望它可以取代 TensorFlow！”。

JAX 结合了 Autograd 和 XLA，是专为高性能机器学习研究打造的产品。

有了新版本的 Autograd，JAX 能够自动对 Python 和 NumPy 的自带函数求导，支持循环、分支、递归、闭包函数求导，而且可以求三阶导数。它支持自动模式反向求导（也就是反向传播）和正向求导，且二者可以任意组合成任何顺序。

JAX 的创新之处在于，它基于 XLA 在 GPU 和 TPU 上编译和运行 NumPy 程序。默认情况下，编译是在底层进行的，库调用能够及时编译和执行。但是 JAX 还允许使用单一函数 API jit将自己的 Python 函数及时编译成经过 XLA 优化的内核。编译和自动求导可以任意组合，因此可以在不脱离 Python 环境的情况下实现复杂算法并获得最优性能。

JAX 最初由 Matt Johnson、Roy Frostig、Dougal Maclaurin 和 Chris Leary 发起，他们均任职于谷歌大脑团队。在 GitHub 的说明文档中，作者明确表示：JAX 目前还只是一个研究项目，不是谷歌的官方产品，因此可能会有一些 bug。从作者的 GitHub 简介来看，这应该是谷歌大脑正在尝试的新项目，在同一个 GitHub 目录下的开源项目还包括 8 月份在业内引起热议的强化学习框架 Dopamine。

以下是 JAX 的简单使用示例。

GitHub 项目传送门：https://github.com/google/JAX

有关具体的安装和简单的入门指导大家可以在 GitHub 中自行查看，在此不做过多赘述。

JAX 库的实现原理

机器学习中的编程是关于函数的表达和转换。转换包括自动微分、加速器编译和自动批处理。像 Python 这样的高级语言非常适合表达函数，但是通常使用者只能应用它们。我们无法访问它们的内部结构，因此无法执行转换。

JAX 可以用于专门化高级 Python+NumPy 函数，并将其转换为可转换的表示形式，然后再提升为 Python 函数。

JAX 通过跟踪专门处理 Python 函数。跟踪一个函数意味着：监视应用于其输入，以产生其输出的所有基本操作，并在有向无环图 (DAG) 中记录这些操作及其之间的数据流。为了执行跟踪，JAX 包装了基本的操作，就像基本的数字内核一样，这样一来，当调用它们时，它们就会将自己添加到执行的操作列表以及输入和输出中。为了跟踪这些原语之间的数据流，跟踪的值被包装在 Tracer 类的实例中。

当 Python 函数被提供给 grad 或 jit 时，它被包装起来以便跟踪并返回。当调用包装的函数时，我们将提供的具体参数抽象到 AbstractValue 类的实例中，将它们框起来用于跟踪跟踪器类的实例，并对它们调用函数。

抽象参数表示一组可能的值，而不是特定的值：例如，jit 将 ndarray 参数抽象为抽象值，这些值表示具有相同形状和数据类型的所有 ndarray。相反，grad 抽象 ndarray 参数来表示底层值的无穷小邻域。通过在这些抽象值上跟踪 Python 函数，我们确保它足够专门化，以便转换是可处理的，并且它仍然足够通用，以便转换后的结果是有用的，并且可能是可重用的。然后将这些转换后的函数提升回 Python 可调用函数，这样就可以根据需要跟踪并再次转换它们。

JAX 跟踪的基本函数大多与 XLA HLO 1:1 对应，并在 lax.py 中定义。这种 1:1 的对应关系使得到 XLA 的大多数转换基本上都很简单，并且确保我们只有一小组原语来覆盖其他转换，比如自动微分。 jax.numpy 层是用纯 Python 编写的，它只是用 LAX 函数 (以及我们已经编写的其他 numpy 函数) 表示 numpy 函数。这使得 jax.numpy 易于延展。

当你使用 jax.numpy 时，底层 LAX 原语是在后台进行 jit 编译的，允许你在加速器上执行每个原语操作的同时编写不受限制的 Python+ numpy 代码。

但是 JAX 可以做更多的事情：你可以在越来越大的函数上使用 jit 来进行端到端编译和优化，而不仅仅是编译和调度到一组固定的单个原语。例如，可以编译整个网络，或者编译整个梯度计算和优化器更新步骤，而不仅仅是编译和调度卷积运算。

折衷之处是，jit 函数必须满足一些额外的专门化需求：因为我们希望编译专门针对形状和数据类型的跟踪，但不是专门针对具体值的跟踪，所以 jit 装饰器下的 Python 代码必须适用于抽象值。如果我们尝试在一个抽象的 x 上求 x >0 的值，结果是一个抽象的值，表示集合{True, False}，所以 Python 分支就像 if x > 0 会引起报错。

有关使用 jit 的更多要求，请参见：https://github.com/google/jax#whats-supported

好消息是，jit 是可选的：JAX 库在后台对单个操作和函数使用 jit，允许编写不受限制的 Python+Numpy，同时仍然使用硬件加速器。但是，当你希望最大化性能时，通常可以在自己的代码中使用 jit 编译和端到端优化更大的函数。

后续计划

目前项目小组还将对以下几项做更多尝试和更新：

1. 完善说明文档
2. 支持 Cloud TPU
3. 支持多 GPU 和多 TPU
4. 支持完整的 NumPy 功能和部分 SciPy 功能
5. 全面支持 vmap
6. 加速
7. 降低 XLA 函数调度开销
8. 线性代数例程（CPU 上的 MKL 和 GPU 上的 MAGMA）
9. 高效自动微分原语cond和while

有关 JAX 库的介绍大致如此，如果你在尝试了 JAX 之后有一些较好的使用心得，欢迎随时向我们投稿，AI 前线十分愿意将你的经验传播给更多开发者。

再次附上 GitHub 链接：https://github.com/google/jax

相关资源：

JAX 论文链接：https://www.sysml.cc/doc/146.pdf

会议推荐

12 月 20-21，AICon将于北京开幕，在这里可以学习来自 Google、微软、BAT、360、京东、美团等 40+AI 落地案例，与国内外一线技术大咖面对面交流。