JAX

JAX 是什么

JAX 把自己描述为一个面向加速器的 Python 数组计算与程序变换库，用于高性能数值计算和大规模机器学习。^[1] 这句话的两半同样重要。数组计算意味着程序处理的基本对象是多维数组，也就是排成网格的数字，对它们的操作是常见的那几类：逐元素运算、矩阵乘法、沿某个轴做归约。程序变换意味着这个库可以接收一个普通的 Python 函数，生成它的一个改写版本——编译过的、批量化了的、可求导的或并行化了的——而不需要用户手工重写这个函数。

JAX 把自己描述为范围很窄的库，专注于高效的数组运算和程序变换。^[1] 它不把自己定位成一套完整的深度学习框架，也不内置现成的层、优化器和训练循环。这种「窄」是有意为之：机器学习需要的那部分由建立在其上的库提供，JAX 则专注于把底层的数组运算和变换做得快、行为可预期。

变换函数

JAX 提供可组合的函数变换，覆盖编译、批处理、自动微分和并行化。^[1] 用白话说，这是四件不同的事。编译把一个 Python 函数整体转成加速器可以执行的代码，这样一长串小的数组操作就不必每一步都经过 Python 解释器。批处理让一个针对单个样本写的函数直接跑在一整叠样本上，不需要显式循环。自动微分生成一个返回原函数导数的函数，训练模型正需要这个。并行化则把计算分摊到多个设备上。

可组合的意思是这些变换可以层层嵌套：批处理过的函数可以再编译，编译过的函数可以再求导，顺序按问题的需要来定。手工写出这样的组合通常需要大量专门代码；这个设计的意义就在于，组合变成了把一个函数套进另一个函数。

一套代码，多种后端

同一份代码可以在多个后端上执行，包括 CPU、GPU 和 TPU。^[1] 这和多数 Python 用户熟悉的数组库不是一回事。NumPy 的运算是即时（eager）、同步执行的，而且只在 CPU 上跑；JAX 的运算既可以即时执行，也可以在 jit 编译之后执行，派发是异步的，并且能在 CPU、GPU 或 TPU 上执行。^[2] 即时执行指调用一次就算一次；用 jit 编译时，函数被追踪一遍，之后重复调用用的是编译好的版本。异步派发指调用在设备算完之前就返回，主机端程序可以接着往下走，真正需要结果时才等待。

面向规模的工具

JAX 带有提前编译（ahead-of-time compilation）、分片（sharding）与自动并行化，以及用 shard_map 做逐设备编程的工具。^[1] 提前编译指编译产物在程序开始运行之前就生成，而不是等到第一次调用。分片指把一个数组切开放到多个设备上，自动并行化让编译器自己推导这次切分需要哪些通信。当自动安排不符合问题需要时，shard_map 提供了一种逐设备显式写代码的方式。

与 NumPy 的比较

用户最常问的就是它和 NumPy 的关系，JAX 的文档也是把它当作性能问题来回答，而不是立场问题。针对 CPU 上单个数组操作的微基准测试，通常对 NumPy 有利，原因是它每次调用的派发开销更低；而 GPU 或 TPU 上的工作负载，以及经过 jit 编译的操作序列，通常对 JAX 有利。^[2] 原因在规模：对一个数组做一次小操作时，调用派发的开销占了大头，NumPy 的路径更短；而在长串操作或加速器上，编译消掉了大部分这类开销，硬件能同时做更多事。

生态与学习路径

因为 JAX 把搭模型这一层留给别人，建立在 JAX 之上的神经网络库包括 Flax、Equinox 和 Keras。^[1] 文档把学习内容分成 6 个层级，从讲数组和 jax.grad 的 JAX 101，一直到讲 jaxpr 语言和 Autodidax 的 JAX 601。^[1] 这条路径从大多数人每天都会用到的部分起步，最后落到变换体系内部更底层的机制。

把这几点放在一起看，JAX 的位置就比较清楚了：它提供的不是一套一站式框架，而是一种可被上层库复用的底座。落在它范围内的东西——数组运算和可组合的变换——被尽量做得快、做得统一；落在范围外的，则交给生态。