『JAX中文文档』JAX快速入门
最新的
原文： https://jax.readthedocs.io/en/latest/notebooks/quickstart.html
JAX快速入门
首先解答一个问题：JAX是什么？
简单的说就是GPU加速、支持自动微分(autodiff)的numpy。众所周知，numpy是Python下的基础数值运算库，得到广泛应用。用Python搞科学计算或机器学习，没人离得开它。但是numpy不支持GPU或其他硬件加速器，也没有对backpropagation的内置支持，再加上Python本身的速度限制，所以很少有人会在生产环境下直接用numpy训练或部署深度学习模型。这也是为什么会出现Theano, TensorFlow, Caffe等深度学习框架的原因。但是numpy有其独特的优势：底层、灵活、调试方便、API稳定且为大家所熟悉（与MATLAB一脉相承），深受研究者的青睐。JAX的主要出发点就是将numpy的以上优势与硬件加速结合。现在已经开源的JAX ( https://github.com/google/jax) 就是通过GPU (CUDA)来实现硬件加速。出自：https://www.zhihu.com/question/306496943/answer/557876584
小宋说：JAX 其实就是一个支持加速器（GPU 和 TPU）的科学计算库(numpy, scipy)和神经网络库(提供relu，sigmoid, conv 等），相较于PyTorch与TensorFlow更加灵活，通用性更佳。这也是笔者推荐学习和做这个翻译工作的原因，带着大家一起去学习掌握这个框架。 由于笔者非英语专业，有些内荣难免翻译有误，欢迎大家批评指正。对于有些笔者不确定的翻译，采用下划线加括号引用原词的方式来补充，例如：自动微分（differentiation）
官方定义：JAX是CPU，GPU和TPU上的NumPy，具有出色的自动差分（differentiation），可用于高性能机器学习研究。
作为更新版本的Autograd，JAX可以自动微分本机Python和NumPy代码。它可以通过Python的大部分功能（包括循环，if，递归和闭包）进行微分，甚至可以采用派生类的派生类。它支持反向模式和正向模式微分，并且两者可以任意顺序组成。
新功能是JAX使用 XLA 在诸如GPU和TPU的加速器上编译和运行您的NumPy代码。默认情况下，编译是在后台进行的，而库调用将得到及时的编译和执行。但是，JAX甚至允许您使用单功能API即时将自己的Python函数编译为XLA优化的内核。编译和自动微分可以任意组合，因此您无需离开Python即可表达复杂的算法并获得最佳性能。
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random乘法矩阵
在以下示例中，我们将生成随机数据。NumPy和JAX之间的一大区别是生成随机数的方式。有关更多详细信息，请参见JAX中的Common Gotchas。
key = random.PRNGKey(0)
x = random.normal(key, (10,))
print(x)[-0.372111 0.2642311 -0.18252774 -0.7368198 -0.44030386 -0.15214427
-0.6713536 -0.59086424 0.73168874 0.56730247]乘以两个大矩阵。
size = 3000
x = random.normal(key, (size, size), dtype=jnp.float32)
%timeit jnp.dot(x, x.T).block_until_ready() # runs on the GPU489 ms ± 3.98 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)我们补充说，block_until_ready因为默认情况下JAX使用异步执行（请参见异步调度）。
JAX NumPy函数可在常规NumPy数组上使用。
import numpy as np
x = np.random.normal(size=(size, size)).astype(np.float32)
%timeit jnp.dot(x, x.T).block_until_ready()488 ms ± 942 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)这样比较慢，因为它每次都必须将数据传输到GPU。您可以使用来确保NDArray由设备内存支持device_put()。
from jax import device_put
x = np.random.normal(size=(size, size)).astype(np.float32)
x = device_put(x)
%timeit jnp.dot(x, x.T).block_until_ready()487 ms ± 9.94 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)的输出device_put()仍然像NDArray一样，但是它仅在需要打印，绘图，保存（printing, plotting, saving）到磁盘，分支等需要它们的值时才将值复制回CPU。的行为device_put()等效于函数，但是速度更快。jit(lambda x: x)
如果您有GPU（或TPU！），这些调用将在加速器上运行，并且可能比在CPU上快得多。
x = np.random.normal(size=(size, size)).astype(np.float32)
%timeit np.dot(x, x.T)235 ms ± 546 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)JAX不仅仅是一个由GPU支持的NumPy。它还带有一些程序转换，这些转换在编写数字代码时很有用。目前，主要有三个：
让我们一一介绍。我们还将最终以有趣的方式编写这些内容。
利用jit()加快功能
JAX在GPU上透明运行（如果没有，则在CPU上运行，而TPU即将推出！）。但是，在上面的示例中，JAX一次将内核分配给GPU一次操作。如果我们有一系列操作，则可以使用@jit装饰器使用XLA一起编译多个操作。让我们尝试一下。
def selu(x, alpha=1.67, lmbda=1.05):
return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
x = random.normal(key, (1000000,))
%timeit selu(x).block_until_ready()4.4 ms ± 107 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)我们可以使用加快速度@jit，它将在第一次selu调用jit-compile并将其之后缓存。
selu_jit = jit(selu)
%timeit selu_jit(x).block_until_ready()860 µs ± 27.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)通过 grad()计算梯度
除了评估数值函数外，我们还希望对其进行转换。一种转变是自动微分。在JAX中，就像在Autograd中一样，您可以使用grad()函数来计算梯度。
def sum_logistic(x):
return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))
x_small = jnp.arange(3.)
derivative_fn = grad(sum_logistic)
print(derivative_fn(x_small))[0.25 0.19661197 0.10499357]让我们以极限微分（finite differences）验证我们的结果是正确的。
def first_finite_differences(f, x):
eps = 1e-3
return jnp.array([(f(x + eps * v) - f(x - eps * v)) / (2 * eps)
for v in jnp.eye(len(x))])
print(first_finite_differences(sum_logistic, x_small))[0.24998187 0.1964569 0.10502338]求解梯度可以通过简单调用grad()。grad()并jit()可以任意混合。在上面的示例中，我们先抖动sum_logistic然后取其派生词。我们继续深入学习实验：
print(grad(jit(grad(jit(grad(sum_logistic)))))(1.0))-0.035325594对于更高级的autodiff，可以将其jax.vjp()用于反向模式矢量雅各比积和jax.jvp()正向模式雅可比矢量积。两者可以彼此任意组合，也可以与其他JAX转换任意组合。这是组合它们以构成有效计算完整的Hessian矩阵的函数的一种方法：
from jax import jacfwd, jacrev
def hessian(fun):
return jit(jacfwd(jacrev(fun)))自动向量化 vmap()
JAX在其API中还有另一种转换，您可能会发现它有用：vmap()向量化映射。它具有沿数组轴映射函数的熟悉语义（ familiar semantics），但不是将循环保留在外部，而是将循环推入函数的原始操作中以提高性能。当与组合时jit()，它的速度可以与手动添加批处理尺寸一样快。
我们将使用一个简单的示例，并使用将矩阵向量乘积提升为矩阵矩阵乘积vmap()。尽管在这种特定情况下很容易手动完成此操作，但是相同的技术可以应用于更复杂的功能。
mat = random.normal(key, (150, 100))
batched_x = random.normal(key, (10, 100))
def apply_matrix(v):
return jnp.dot(mat, v)给定诸如之类的功能apply_matrix，我们可以在Python中循环执行批处理维度，但是这样做的性能通常很差。
def naively_batched_apply_matrix(v_batched):
return jnp.stack([apply_matrix(v) for v in v_batched])
print('Naively batched')
%timeit naively_batched_apply_matrix(batched_x).block_until_ready()Naively batched4.43 ms ± 9.91 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)我们知道如何手动批处理此操作。在这种情况下，jnp.dot透明地处理额外的批次尺寸。
@jit
def batched_apply_matrix(v_batched):
return jnp.dot(v_batched, mat.T)
print('Manually batched')
%timeit batched_apply_matrix(batched_x).block_until_ready()Manually batched51.9 µs ± 1.72 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)但是，假设没有批处理支持，我们的功能更加复杂。我们可以用来vmap()自动添加批处理支持。
@jit
def vmap_batched_apply_matrix(v_batched):
return vmap(apply_matrix)(v_batched)
print('Auto-vectorized with vmap')
%timeit vmap_batched_apply_matrix(batched_x).block_until_ready()Auto-vectorized with vmap79.7 µs ± 249 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)当然，vmap()可以与任意组成jit()，grad()和任何其它JAX变换。
这只是JAX可以做的事情。我们很高兴看到您的操作！