brew search python@3.10
brew install python@3.10
1
2

Pouring sqlite-3.38.5.arm64_monterey.bottle.tar.gz
Error: No such file or directory @ rb_sysopen - /Users/xzh/Library/Caches/Homebrew/downloads/062e09dc048eab6bed4b64a9ce0533b08d65775640f901d27e24fd4c1ae640d7--sqlite-3.38.5.arm64_monterey.bottle.tar.gz
1
2

#Setting PATH for Python 3.10 
export PATH=${PATH}:/Library/Frameworks/Python.framework/Versions/3.10/bin
alias python="/Library/Frameworks/Python.framework/Versions/3.10/bin/python3.10"
export PATH=${PATH}:/Library/Frameworks/Python.framework/Versions/3.10/bin
alias pip="/Library/Frameworks/Python.framework/Versions/3.10/bin/pip3"
export PATH="$PATH:/usr/local/bin/python3.10"
1
2
3
4
5
6

Python 3.10.5 (v3.10.5:f377153967, Jun  6 2022, 12:36:10) [Clang 13.0.0 (clang-1300.0.29.30)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> 
1
2
3

curl https://bootstrap.pypa.io/get-pip.py | python3
1

pip._vendor.urllib3.exceptions.ReadTimeoutError: HTTPSConnectionPool(host='files.pythonhosted.org', port=443): Read timed out.
1

mkdir -p ~/.pip
cat > ~/.pip/pip.conf<

• 1
• 2
• 3
• 4
• 5
• 6
• 7

然后，执行安装如果有下面的提示，则说明镜像源已被替换。

Looking in indexes: https://mirrors.aliyun.com/pypi/simple/
1

2.3 JAX基本使用

2.3.1 JAX插件安装

经过前面的介绍，我们知道，jax其实就是一个函数库，所以我们使用之前，需要先安装一下jax插件，安装时需要使用pip命令安装，命令如下：

pip install --upgrade pip
pip install --upgrade "jax[cpu]"
1
2

关于如何安装，大家可以参考下官方文档的介绍。安装成功之后，会给出成功的提示，如下图。

除了CPU版本外，JAX还支持GPU和TPU，安装的命令如下：

//GPU
pip install --upgrade pip
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
//TPU
pip install --upgrade pip
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
1
2
3
4
5
6

2.3.2 官方示例工程

JAX 的定位是科学计算（Scientific Computing）和函数转换（Function Transformations），具有除训练深度学习模型以外的一系列能力，具体包括：

• 即时编译（Just-in-Time Compilation）
• 自动并行化（Automatic Parallelization）
• 自动向量化（Automatic Vectorization）
• 自动微分（Automatic Differentiation）

为了方便学习，我们最好下载下官网的示例工程并运行，示例工程代码结构如下：

2.3.3 在线编程平台jupyter

Jupyter 是一款开放性的代码编写软件，最大的特点是能够实时运行代码，查看输出效果。同时该软件集成了多种插件，能够实现非常复杂的功能。对于初学者而言，因为Jupyter具有良好的交互性，方便看到每一行、每一个代码块的输出结果，Jupyter成为众多Python初学者的青睐。

同时，Jupyter提供了用 JupyterLab 和 Jupyter Notebook 等交互式编写软件的技术方式，能够更好的帮助开发者编写、运行代码。其中，安装JupyterLab命令如下：

pip3 install jupyterlab 
//启动命令
jupyter-lab
1
2
3

安装完成之后，在使用jupyter-lab即可启动。如果是安装Jupyter Notebook，那么安装的命令如下：

pip3 install notebook 
//启动命令
jupyter notebook
1
2
3

启动成功之后，会自动打开http://localhost:8888/tree页面，如下图。

然后，我们点击右上角的【新建】按钮新建一个运行面板，如下图。

接下来，我们就可以在上面运行一些函数。当然，Jupyter还提供了在线编辑运行平台，可以帮助开发者快速的体验Jupyter的魅力，目前支持主流的编程语言和技术。

如果要运行项目，可以使用【Ctrol】+回车即可得到运行结果。比如生成随机数据：

key = random.PRNGKey(0)
x = random.normal(key, (10,))
print(x)
1
2
3

当我们【Ctrol】+回车运行项目时，得到的结果如下：

[-0.3721109   0.26423115 -0.18252768 -0.7368197  -0.44030377 -0.1521442
 -0.67135346 -0.5908641   0.73168886  0.5673026 ]
1
2

其中，比较常用的快捷键有如下一些：

• Tab : 代码补全或缩进
• Shift-Tab : 提示
• Ctrl-A : 全选
• Ctrl-Z : 复原
• Ctrl-Shift-Z : 再做
• Ctrl-Y : 再做
• Ctrl-Backspace : 删除前面一个字
• Ctrl-Delete : 删除后面一个字
• Esc : 进入命令模式
• Ctrl-M : 进入命令模式
• Shift-Enter : 运行本单元，选中下一单元
• Ctrl-Enter : 运行本单元
• Alt-Enter : 运行本单元，在下面插入一单元
• Ctrl-Shift-- : 分割单元
• Ctrl-Shift-Subtract : 分割单元
• Ctrl-S : 文件存盘
• Shift : 忽略
  其他常用的快捷键，可以通过帮助选项来进行查看，如下图。
  

2.3.4 基本使用

随机函数

和编写其他python的语法一样，使用Jax之前需要导入相关的函数包。比如，我们使用NumPy执行一些基准测试，比如：

import jax
import jax.numpy as jnp
from jax import random
from jax import grad, jit
import numpy as np
key = random.PRNGKey(0)
1
2
3
4
5
6

当然，我们也可以 import jax.numpy as jnp 并将代码中的所有 np 替换为 jnp 。与NumPy 代码风格不同，在JAX 代码中，可以直接使用import方式导入并直接使用。可以看到，JAX 中随机数的生成方式与 NumPy 不同。JAX需要创建一个 jax.random.PRNGKey 。

矩阵乘法

我们在 Google Colab 上做一个简单的基准测试，这样我们就可以轻松访问 GPU 和 TPU。我们首先初始化一个包含 25M 元素的随机矩阵，然后将其乘以它的转置，使用针对 CPU 优化的 NumPy，矩阵乘法平均需要408 ms ± 35.9 ms。

size = 5000
x = np.random.normal(size=(size, size)).astype(np.float32)
%timeit np.dot(x, x.T)

# 408 ms ± 35.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1
2
3
4
5

在 CPU 上使用 JAX 执行相同的操作平均需要大约 716 ms ± 13.7 ms。

size = 5000
x = random.normal(key, (size, size), dtype=jnp.float32)
%timeit jnp.dot(x, x.T).block_until_ready()

# 716 ms ± 13.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1
2
3
4
5

在 CPU 上运行时，JAX 通常比 NumPy 慢，因为 NumPy 已针对CPU进行了非常多的优化。但是，当使用加速器时这种情况会发生变化，所以让我们尝试使用 GPU 进行矩阵乘法。

size = 5000
x = random.normal(key, (size, size), dtype=jnp.float32)
%time x_jax = jax.device_put(x)  
%time jnp.dot(x_jax, x_jax.T).block_until_ready()  
%timeit jnp.dot(x_jax, x_jax.T).block_until_ready() 

# CPU times: user 50 µs, sys: 1 µs, total: 51 µs
# Wall time: 53.9 µs
# CPU times: user 5.14 s, sys: 44.9 ms, total: 5.19 s
# Wall time: 732 ms
# 725 ms ± 25.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1
2
3
4
5
6
7
8
9
10
11

接下来，让我们使用 TPU 来进行矩阵乘法。

size = 5000
x = random.normal(key, (size, size), dtype=jnp.float32)
%time x_jax = jax.device_put(x)  
%time jnp.dot(x_jax, x_jax.T).block_until_ready()  
%timeit jnp.dot(x_jax, x_jax.T).block_until_ready() 

# CPU times: user 54 µs, sys: 916 µs, total: 970 µs
# Wall time: 973 µs
# CPU times: user 5.25 s, sys: 34.2 ms, total: 5.28 s
# Wall time: 709 ms
# 715 ms ± 3.95 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1
2
3
4
5
6
7
8
9
10
11

可以看到，忽略设备传输时间和编译时间，每个矩阵乘法平均需要715 ms ± 3.95 毫秒，与GPU 相比，TPU快了差不多4倍。需要说明的是，当乘以不同大小的矩阵时，获得相同的加速效果也不同：相乘的矩阵越大，GPU可以优化操作的越多，加速也越大。

jit()

JAX在GPU上是透明运行的。但是，在上面的示例中，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()

#1.64 ms ± 91 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1
2
3
4
5
6
7

可以使用加快速度@jit，它将在第一次selu调用jit-compile并将其之后缓存。

def selu(x, alpha=1.67, lmbda=1.05):
  return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

selu_jit = jit(selu)
%timeit selu_jit(x).block_until_ready()

#455 µs ± 151 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
1
2
3
4
5
6
7

可以看到，使用jit装饰器后，运行效率明显提高。

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]
1
2
3
4
5
6
7
8

接下来，让我们使用极限微分来验证我们的结果是否正确。

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]
1
2
3
4
5
6
7
8

求解梯度可以通过简单调用grad()。grad()并jit()可以任意混合。在上面的示例中，我们先抖动sum_logistic然后取其派生词。

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(grad(jit(grad(jit(grad(sum_logistic)))))(1.0))

//-0.035325598
1
2
3
4
5
6
7
8

vmap()

JAX在其API中还有另一种转换，那就是vmap()向量化映射。它具有沿数组轴映射函数的熟悉语义，但不是将循环保留在外部，而是将循环推入函数的原始操作中以提高性能。当与组合时jit()，它的速度可以与手动添加批处理尺寸一样快。

mat = random.normal(key, (150, 100))
batched_x = random.normal(key, (10, 100))
 
def apply_matrix(v):
  return jnp.dot(mat, v)

print(apply_matrix(100))
1
2
3
4
5
6
7

给定功能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 batched
# 433 µs ± 2.02 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1
2
3
4
5
6
7
8

如果，我们使用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 vmap
# 13.5 µs ± 19.4 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
1
2
3
4
5
6
7
8
9

事实上，vmap()可以与任意组成jit()，grad()和任何其它JAX变换。当然，JAX的函数还有很多，大家可以查看官方资料进行学习。

三、XLA架构

XLA 是 JAX（和其他库，例如 TensorFlow，TPU的Pytorch）使用的线性代数的编译器，它通过创建自定义优化内核来保证最快的在程序中运行线性代数运算。XLA 最大的好处是可以让我们在应用中自定义内核，该部分使用线性代数运算，以便它可以进行最多的优化。

在TensorFlow中，XLA给TensorFlow带来了如下提升：

• 提高执行速度。编译子计算图以减少短暂运算的执行时间，从而消除 TensorFlow 运行时的开销；融合流水线运算以降低内存开销；并针对已知张量形状执行专门优化以支持更积极的常量传播。
• 提高内存使用率。分析和安排内存使用量，原则上需要消除许多中间存储缓冲区。
• 降低对自定义运算的依赖。通过提高自动融合的低级运算的性能，使之达到手动融合的自定义运算的性能水平，从而消除对多种自定义运算的需求。
• 减少移动资源占用量。通过提前编译子计算图并发出可以直接链接到其他应用的对象/头文件对，消除 TensorFlow 运行时。这样，移动推断的资源占用量可降低几个数量级。
• 提高便携性。使针对新颖硬件编写新后端的工作变得相对容易，在新硬件上运行时，大部分 TensorFlow 程序都能够以未经修改的方式运行。与针对新硬件专门设计各个整体运算的方式相比，这种模式不必重新编写 TensorFlow 程序即可有效利用这些运算。

不过，XLA 最重要的优化是融合，即可以在同一个内核中进行多个线性代数运算，将中间输出保存到 GPU 寄存器中，而不将它们具体化到内存中。这可以显着增加我们的“计算强度”，即所做的工作量与负载和存储数量的比例。融合还可以让我们完全省略仅在内存中shuffle 的操作（例如reshape）。

下面我们看看如何使用 XLA 和 jax.jit 手动触发 JIT 编译。

使用 jax.jit 进行即时编译

这里有一些新的基准来测试 jax.jit 的性能。我们定义了两个实现 SELU（Scaled Exponential Linear Unit）的函数：一个使用 NumPy，一个使用 JAX。

def selu_np(x, alpha=1.67, lmbda=1.05):
  return lmbda * np.where(x > 0, x, alpha * np.exp(x) - alpha)
def selu_jax(x, alpha=1.67, lmbda=1.05):
  return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

x = np.random.normal(size=(1000000,)).astype(np.float32)
%timeit selu_np(x)

# 7.56 ms ± 18.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
1
2
3
4
5
6
7
8
9

可以看到，NumPy平均需要 7.6 毫秒。接下来，让我们在 CPU 上使用 JAX运行，如下。

def selu_np(x, alpha=1.67, lmbda=1.05):
  return lmbda * np.where(x > 0, x, alpha * np.exp(x) - alpha)
def selu_jax(x, alpha=1.67, lmbda=1.05):
  return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

x = random.normal(key, (1000000,))
%time selu_jax(x).block_until_ready()  
%timeit selu_jax(x).block_until_ready()

# CPU times: user 5.27 ms, sys: 2.7 ms, total: 7.97 ms
# Wall time: 3.57 ms
# 1.7 ms ± 94.6 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1
2
3
4
5
6
7
8
9
10
11
12

可以看到，这种情况下，JAX明显要比 NumPy 快。下一个测试是在 GPU 上使用 JAX。

x = random.normal(key, (1000000,))
%time x_jax = jax.device_put(x)  
%time selu_jax(x_jax).block_until_ready()  
%timeit selu_jax(x_jax).block_until_ready() 

# CPU times: user 54 µs, sys: 39 µs, total: 93 µs
# Wall time: 96.1 µs
# CPU times: user 2.27 ms, sys: 1.3 ms, total: 3.57 ms
# Wall time: 1.71 ms
# 1.63 ms ± 45.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1
2
3
4
5
6
7
8
9
10

可以看到，函数运行时间为1.63毫秒。下面我们用 jax.jit 测试它，触发 JIT 编译器使用 XLA 将 SELU 函数编译到优化的 GPU 内核中，同时优化函数内部的所有操作。

x = random.normal(key, (1000000,))
selu_jax_jit = jit(selu_jax)
%time x_jax = jax.device_put(x)  
%time selu_jax_jit(x_jax).block_until_ready()  
%timeit selu_jax_jit(x_jax).block_until_ready() 

# CPU times: user 114 µs, sys: 305 µs, total: 419 µs
# Wall time: 426 µs
# CPU times: user 30.2 ms, sys: 7.86 ms, total: 38.1 ms
# Wall time: 36.5 ms
# 361 µs ± 10.1 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
1
2
3
4
5
6
7
8
9
10
11

可以看到，使用编译内核，函数运行时间为0.36毫秒。之所以能带来如此大的性能提升，是因为使用 JIT 编译避免从 GPU 寄存器中移动数据，从未带来了非常大的加速。一般来说在不同类型的内存之间移动数据与代码执行相比非常慢，因此在实际使用时应该尽量避免。

将 SELU 函数应用于不同大小的向量时，您可能会获得不同的结果。矢量越大，加速器越能优化操作，加速也越大。除了执行 selu_jax_jit = jit(selu_jax) 之外，还可以使用 @jit 装饰器对函数进行 JIT 编译，如下所示。

@jit
def selu_jax_jit(x, alpha=1.67, lmbda=1.05):
  return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
1
2
3

JIT 编译可以加速，为什么我们不能全部都这样做呢？因为并非所有代码都可以 JIT 编译，JIT要求数组形状是静态的并且在编译时已知。另外就是引入jax.jit 也会带来一些开销。因此通常只有编译的函数比较复杂并且需要多次运行才能节省时间。

参考链接：
https://jax.readthedocs.io/en/latest/installation.html
https://github.com/google/jax