JAX快速上手：从NumPy到GPU加速的Python高性能计算库入门教程

摘要：NumPy作为Python数值计算领域的基础框架，凭借其强大的N维数组结构和丰富的函数生态系统，成为科学家、工程师和数据分析师的核心工具。然而，随着计算需求的快速增长，特别是在机器学习和大规模科学模拟领域，NumPy基于CPU的执行模式以及缺乏内置自动微分功能

NumPy作为Python数值计算领域的基础框架，凭借其强大的N维数组结构和丰富的函数生态系统，成为科学家、工程师和数据分析师的核心工具。然而，随着计算需求的快速增长，特别是在机器学习和大规模科学模拟领域，NumPy基于CPU的执行模式以及缺乏内置自动微分功能的限制愈发明显。

JAX正是为了解决这些问题而设计的。作为Google Research开发的数值计算库，JAX致力于将NumPy引入现代硬件加速器和基于梯度优化的计算范式中。需要说明的是，JAX目前仍是一个研究项目，而非Google的官方产品，因此在使用过程中可能遇到一些不稳定因素和潜在问题。

NumPy的核心特性与局限性

对于从事科学计算的Python开发者而言，NumPy几乎是必备工具。其核心ndarray对象为高效存储和操作密集数值数组提供了基础，配合大量经过优化的数学函数库（通常通过C、C++或Fortran代码实现），在处理数组运算时远超纯Python循环的性能。

NumPy构成了整个科学Python生态系统的基石，包括SciPy、Pandas、Scikit-learn、Matplotlib等重要库都建立在其之上。然而，其设计主要面向CPU执行环境，且本身不具备自动梯度计算能力。

梯度计算的重要性不容忽视。梯度衡量函数输出相对于输入的变化率，指示函数值增减最快的方向。这一信息对优化算法至关重要，特别是在机器学习领域，梯度指导模型参数的调整过程，以在训练期间最小化损失函数。高效的梯度计算能力使模型能够从数据中学习复杂的模式。

JAX技术架构与核心功能

JAX是Google开发的高性能Python数值计算库，将类似NumPy的API与自动微分（autodiff）和加速硬件执行（GPU/TPU）相结合。

JAX的主要技术特性体现在以下几个方面：jax.numpy模块提供了NumPy的直接替代方案，保持相同的API接口但支持GPU/TPU加速；jax.grad模块实现函数的自动微分功能，类似于TensorFlow或PyTorch的梯度计算机制；jax.jit模块通过加速线性代数（XLA）库提供即时编译功能，实现极高的执行效率；jax.vmap和jax.pmap模块支持自动向量化和并行化处理；此外JAX在GPU/TPU上无缝运行，无需修改现有代码。

JAX的应用场景分析

JAX在以下场景中表现出显著优势：当需要通过GPU或TPU显著加速类似NumPy的计算任务时；当需要自动计算数值Python函数的梯度以进行优化时（如机器学习、物理模拟等）；当需要通过JIT编译关键Python代码段以获得进一步加速时；当需要轻松实现函数向量化以处理数据批次或在多个加速器设备上并行化计算时。

在选择JAX替代NumPy之前，需要了解两个库之间的关键差异。虽然jax.numpy在很大程度上模仿NumPy API，但在以下方面存在显著区别：

在执行后端和编译方面，NumPy在CPU上采用即时执行模式，通常使用预编译的C、C++或Fortran扩展以及优化的线性代数库（如OpenBLAS）。相比之下JAX使用XLA编译器将代码转换为针对CPU、GPU或TPU优化的机器代码，支持通过jax.jit进行即时编译，并通常采用异步分派方式。

在执行模型方面，NumPy操作通常同步执行，Python解释器等待操作完成后继续执行。而JAX操作异步分派到加速器，Python代码可能在计算进行时继续运行。因此，通常需要使用result.block_until_ready来确保准确计时或在其他地方使用结果之前确保结果可用。

在数据可变性方面，NumPy数组（ndarray）是可变的，允许就地修改元素。JAX数组则是不可变的，不允许就地更新。这种函数式编程方法确保JAX的转换功能能够可靠工作而不产生副作用，更新操作需要使用索引更新语法创建新数组。

在随机数生成方面，NumPy使用全局随机数生成器状态，这在并行或转换代码中可能影响可重现性。JAX则需要显式处理随机密钥，必须手动管理和分割密钥以确保随机性的可重现性。

在API覆盖范围方面，NumPy提供涵盖数值计算各个方面的全面API，而JAX覆盖最常见NumPy API的大部分子集且在不断扩展，但并非100%的直接替代品。一些不常见的函数、特定数据类型（如对象数组）或特定行为可能存在差异或缺失。

环境配置与安装

本文的代码示例基于WSL2 Ubuntu for Windows开发环境。对于拥有Nvidia GPU的系统，可以充分利用GPU加速功能。即使没有GPU，JAX仍能在CPU上提供相比NumPy更好的性能。对于不同的GPU品牌或配置，建议参考官方文档获取详细的安装说明。

首先创建专用的开发环境，这里使用conda进行环境管理：

conda create -n jax_test python=3.13 -y

激活环境并安装必要的库。需要注意的是，对于NVIDIA GPU，需要确保安装了适当的NVIDIA驱动程序和CUDA环境（如CUDA 11或CUDA 12）：

conda activate jax_test
pip install jupyter numpy "jax[cuda12]" matplotlib pillow

JAX的安装过程可能较为耗时，完成后可以启动Jupyter notebook。如果浏览器未自动打开，可以从命令行输出中找到类似以下格式的URL：

http://127.0.0.1:8888/tree?token=3b9f7bd07b6966b41b68e2350721b2d0b6f388d248cc69d

1、熟悉API与JIT编译优化

第一个案例展示了JAX如何通过即时编译技术提升NumPy的性能。这个案例实现了应用于10,000 x 10,000数组的SELU（Scaled Exponential Linear Unit）函数。SELU是自归一化神经网络中广泛使用的激活函数，其数学定义如下图所示：

该函数通过np.where（NumPy）或jnp.where（JAX）实现，根据输入值x的正负性选择不同的计算公式。

代码实现包含三种版本：selu_numpy(x)为标准NumPy实现；selu_jax(x)为JAX版本，代码结构相同但使用JAX数组；selu_jax_jit(x)在JAX版本基础上添加@jax.jit装饰器以启用编译优化。

测试使用一个10,000 x 10,000的随机数数组作为输入数据，分别测量各实现的执行时间。其中，selu_numpy在CPU上直接运行；selu_jax在没有JIT优化的情况下运行，由于需要解释执行而相对较慢；selu_jax_jit首次运行时包含编译时间，但在后续运行中重用编译后的函数，执行速度显著提升。

需要特别注意的是，由于JAX采用异步执行模式，需要使用block_until_ready等待操作完成以获得准确的性能测量结果。首次JIT运行包含编译开销，而第二次运行则直接使用缓存的编译结果，实现显著的性能提升。

import numpy as np
import jax
import jax.numpy as jnp
from timeit import default_timer as timer

# 为SELU定义常数
alpha = 1.6732632423543772848170429916717
scale = 1.0507009873554804934193349852946

# --- NumPy版本 ---
def selu_numpy(x):
return scale * np.where(x > 0, x, alpha * np.exp(x) - alpha)

# --- JAX版本 ---
def selu_jax(x):
return scale * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

# --- JIT编译的JAX版本 ---
# 应用@jax.jit装饰器
@jax.jit
def selu_jax_jit(x):

# 生成测试数据
x_np = np.random.rand(10000, 10000).astype(np.float32)
# 使用JAX的随机数生成（需要显式密钥管理）
key = jax.random.PRNGKey(0)
x_jax = jax.random.normal(key, (10000, 10000), dtype=jnp.float32)

print("执行性能基准测试...")

# --- 性能基准测试 ---

# NumPy基准测试
start = timer
result_np = selu_numpy(x_np)
# NumPy在CPU上同步执行，无需等待
print(f"NumPy执行时间: {timer-start:.6f} 秒")

# JAX（无JIT）基准测试
start = timer
result_jax = selu_jax(x_jax)
result_jax.block_until_ready # 关键步骤：等待JAX计算完成
print(f"JAX (无JIT)执行时间: {timer-start:.6f} 秒")

# JAX（JIT）首次运行（包含编译时间）
start = timer
result_jax_jit = selu_jax_jit(x_jax)
result_jax_jit.block_until_ready
print(f"JAX (JIT)首次运行时间（含编译）: {timer-start:.6f} 秒")

# JAX（JIT）第二次运行（使用缓存编译结果）
start = timer
result_jax_jit_2 = selu_jax_jit(x_jax)
result_jax_jit_2.block_until_ready
print(f"JAX (JIT)第二次运行时间: {timer-start:.6f} 秒")

# 验证计算结果的一致性
print(np.allclose(selu_numpy(np.array(x_jax)), result_jax_jit_2, atol=1e-6))

测试结果显示：

执行性能基准测试...
NumPy执行时间: 0.357104 秒
JAX (无JIT)执行时间: 0.108734 秒
JAX (JIT)首次运行时间（含编译）: 0.026956 秒
JAX (JIT)第二次运行时间: 0.002400 秒
True

结果表明，JAX的JIT编译版本在第二次运行时相比NumPy实现了超过100倍的性能提升，即使不使用JIT优化，JAX的性能也比NumPy提升约3倍。

2、自动微分功能演示

下面这个案例将展示了JAX强大的自动微分能力，这是其区别于NumPy的核心特性之一。

import jax
import jax.numpy as jnp

# 使用jax.numpy定义目标函数
def cubic_sum(x):
return jnp.sum(x**3)

# 通过jax.grad获取梯度函数
grad_cubic_sum = jax.grad(cubic_sum)

# 创建输入数据
x_input = jnp.arange(1.0, 5.0)

# 计算梯度
gradient = grad_cubic_sum(x_input)
print(f"\n--- 自动微分示例 ---")
print(f"原始函数输入: {x_input}")
print(f"函数输出 f(x): {cubic_sum(x_input)}")
print(f"梯度 df/dx: {gradient}")

执行结果：

--- 自动微分示例 ---
原始函数输入: [1. 2. 3. 4.]
函数输出 f(x): 100.0
梯度 df/dx: [ 3. 12. 27. 48.]

在这个示例中，目标函数为x³，其导数为3x²。对于输入序列[1, 2, 3, 4]，相应的梯度计算为：

3 × 1² = 3

3 × 2² = 12

3 × 3² = 27

3 × 4² = 48

这个简单的例子展示了JAX自动微分的强大功能，无需手动推导和实现梯度计算，JAX能够自动处理复杂函数的微分运算。

3、向量化矩阵乘法性能分析

我们还通过向量化操作展示了JAX的性能优势，具体实现了将一个10,000元素矩阵与128个10,000元素向量批次相乘的计算任务。

import numpy as np
import jax
import jax.numpy as jnp
from timeit import default_timer as timer

# --- 单个数据点的基础函数 ---
def mat_vec_product(matrix, vector):
"""计算矩阵-向量乘积"""
return jnp.dot(matrix, vector)

# --- 使用vmap创建批处理版本 ---
# 目标是将mat_vec_product应用到批次中的每个向量
# 矩阵对批次中的所有向量保持不变
# in_axes=(None, 0)的含义：
# None: 不对第一个参数(matrix)进行映射，采用广播方式
# 0: 对第二个参数的第0轴(向量批次)进行映射
batched_mat_vec = jax.vmap(mat_vec_product, in_axes=(None, 0))

# --- JIT编译的向量化函数 ---
@jax.jit
def batched_mat_vec_jit(matrix, vectors):
"""JIT编译的批处理矩阵-向量乘法"""
return jax.vmap(mat_vec_product, in_axes=(None, 0))(matrix, vectors)

# --- 数据配置 ---
matrix_size = 10000
vector_size = 10000
batch_size = 128
dtype = jnp.float32

# JAX随机数生成需要显式密钥管理
key = jax.random.PRNGKey(0)
key, subkey1, subkey2 = jax.random.split(key, 3)

# 生成测试数据
matrix_jax = jax.random.normal(subkey1, (matrix_size, vector_size), dtype=dtype)
vectors_jax = jax.random.normal(subkey2, (batch_size, vector_size), dtype=dtype)

# 转换为NumPy格式以进行对比测试
matrix_np = np.array(matrix_jax)
vectors_np = np.array(vectors_jax)

print(f"\n--- vmap性能基准测试 (矩阵: {matrix_size}x{vector_size}, 批次大小: {batch_size}) ---")
print(f"可用JAX设备: {jax.devices}")

# --- 性能基准测试 ---

# NumPy方法1：Python循环实现（仅作演示，通常性能较差）
start_np_loop = timer
output_np_loop = np.array([np.dot(matrix_np, v) for v in vectors_np])
end_np_loop = timer
print(f"NumPy (Python循环)执行时间: {end_np_loop - start_np_loop:.6f} 秒")

# NumPy方法2：矩阵乘法与转置操作（高效实现方式）
start_np_matmul = timer
# 矩阵乘法要求vectors_np的形状为(vector_size, batch_size)
output_np_matmul = (matrix_np @ vectors_np.T).T
end_np_matmul = timer
print(f"NumPy (矩阵乘法@)执行时间: {end_np_matmul - start_np_matmul:.6f} 秒")

# JAX vmap（无JIT优化）
start_jax_vmap = timer
output_jax_vmap = batched_mat_vec(matrix_jax, vectors_jax)
output_jax_vmap.block_until_ready
end_jax_vmap = timer
print(f"JAX (vmap, 无JIT)执行时间: {end_jax_vmap - start_jax_vmap:.6f} 秒")

# JAX vmap（JIT优化）首次运行（包含编译开销）
start_jax_vmap_jit_compile = timer
output_jax_vmap_jit_compile = batched_mat_vec_jit(matrix_jax, vectors_jax)
output_jax_vmap_jit_compile.block_until_ready
end_jax_vmap_jit_compile = timer
print(f"JAX (vmap+JIT)首次运行时间（含编译）: {end_jax_vmap_jit_compile - start_jax_vmap_jit_compile:.6f} 秒")

# JAX vmap（JIT优化）第二次运行
start_jax_vmap_jit = timer
output_jax_vmap_jit = batched_mat_vec_jit(matrix_jax, vectors_jax)
output_jax_vmap_jit.block_until_ready
end_jax_vmap_jit = timer
print(f"JAX (vmap+JIT)第二次运行时间: {end_jax_vmap_jit - start_jax_vmap_jit:.6f} 秒")

性能测试结果：

--- vmap性能基准测试 (矩阵: 10000x10000, 批次大小: 128) ---
可用JAX设备: [CudaDevice(id=0)]
NumPy (Python循环)执行时间: 1.129315 秒
NumPy (矩阵乘法@)执行时间: 0.029319 秒
JAX (vmap, 无JIT)执行时间: 0.901569 秒
JAX (vmap+JIT)首次运行时间（含编译）: 0.539354 秒
JAX (vmap+JIT)第二次运行时间: 0.001776 秒

虽然首次JIT编译需要相对较长的时间，但后续运行的性能提升极为显著，体现了JAX编译优化的强大效果。

4、图像卷积处理实现

最后一个案例以图像处理中的高斯模糊为例，展示了JAX在实际应用场景中的性能表现。卷积是图像处理的基础操作，广泛应用于模糊、锐化和边缘检测等任务。该操作涉及在图像上滑动小矩阵（卷积核），并在每个位置计算核下像素的加权和。本案例通过数组切片和逐元素操作实现高斯模糊的基本版本，以评估jax.jit对此类操作序列的优化效果。

输入图像示例：

原始图像由Yury Taranik提供

import numpy as np
import jax
import jax.numpy as jnp
from timeit import default_timer as timer
from PIL import Image
import matplotlib.pyplot as plt
import os

# --- 配置参数 ---
image_path = "/mnt/d/images/taj_mahal.png"
kernel_size = 9 # 增大卷积核以获得更明显的模糊效果
sigma = 2.5
dtype = jnp.float32

# --- 图像文件检查 ---
if not os.path.exists(image_path):
print(f"错误：在路径 '{image_path}' 未找到图像文件")
print("请更新脚本中的 'image_path' 变量")
exit

# --- 图像加载和预处理 ---
print(f"从以下路径加载图像: {image_path}")
try:
# 打开图像，转换为灰度模式，然后转换为NumPy数组
with Image.open(image_path) as img:
image_np_uint8 = np.array(img.convert('L'))

# 归一化为0.0到1.0范围的float32类型
image_np = image_np_uint8.astype(np.float32) / 255.0
image_jax = jnp.array(image_np)

image_size_h, image_size_w = image_np.shape
print(f"图像加载成功 ({image_size_h}x{image_size_w})")

except Exception as e:
print(f"错误：无法加载或处理图像 '{image_path}'。错误信息: {e}")
exit

# --- 高斯卷积核生成 ---
def gaussian_kernel(size, sigma=1.0):
"""生成2D高斯卷积核"""
ax = jnp.arange(-size // 2 + 1., size // 2 + 1.)
xx, yy = jnp.meshgrid(ax, ax)
kernel = jnp.exp(-(xx**2 + yy**2) / (2. * sigma**2))
return (kernel / jnp.sum(kernel)).astype(dtype)

# --- 手动卷积实现 ---
def convolve_2d_manual(image, kernel):
"""使用基本数组操作实现2D卷积"""
im_h, im_w = image.shape
ker_h, ker_w = kernel.shape
pad_h, pad_w = ker_h // 2, ker_w // 2
padded_image = jnp.pad(image, ((pad_h, pad_h), (pad_w, pad_w)), mode='edge')
output = jnp.zeros_like(image)
for i in range(ker_h):
for j in range(ker_w):
# 使用dynamic_slice确保与JIT兼容
image_slice = jax.lax.dynamic_slice(padded_image, (i, j), (im_h, im_w))
output += kernel[i, j] * image_slice
return output

# --- JIT编译版本 ---
@jax.jit
def convolve_2d_manual_jit(image, kernel):
"""JIT编译的2D卷积实现"""

# 当切片大小动态变化时，使用jax.lax.dynamic_slice确保JIT兼容性

return output

# --- NumPy对照实现 ---
def convolve_2d_manual_np(image, kernel):
"""NumPy版本的2D卷积实现"""

padded_image = np.pad(image, ((pad_h, pad_h), (pad_w, pad_w)), mode='edge')
output = np.zeros_like(image)

image_slice = padded_image[i:i + im_h, j:j + im_w]

return output

# --- 卷积核准备 ---
kernel_jax = gaussian_kernel(kernel_size, sigma=sigma)
kernel_np = np.array(kernel_jax)

print(f"\n--- 卷积性能基准测试 (图像: {image_size_h}x{image_size_w}, 卷积核: {kernel_size}x{kernel_size}) ---")
print(f"可用JAX设备: {jax.devices}")

# --- 性能基准测试 ---

# NumPy（CPU）实现
start_np = timer
output_np = convolve_2d_manual_np(image_np, kernel_np)
end_np = timer
print(f"NumPy (手动卷积)执行时间: {end_np - start_np:.6f} 秒")

# JAX（无JIT）实现
start_jax = timer
output_jax = convolve_2d_manual(image_jax, kernel_jax)
output_jax.block_until_ready
end_jax = timer
print(f"JAX (无JIT, 手动卷积)执行时间: {end_jax - start_jax:.6f} 秒")

# JAX（JIT）首次运行（包含编译时间）
start_jax_compile = timer
output_jax_jit_compile = convolve_2d_manual_jit(image_jax, kernel_jax)
output_jax_jit_compile.block_until_ready
end_jax_compile = timer
print(f"JAX (JIT, 手动卷积)首次运行时间（含编译）: {end_jax_compile - start_jax_compile:.6f} 秒")

# JAX（JIT）第二次运行
start_jax_jit = timer
output_jax_jit = convolve_2d_manual_jit(image_jax, kernel_jax)
output_jax_jit.block_until_ready
end_jax_jit = timer
print(f"JAX (JIT, 手动卷积)第二次运行时间: {end_jax_jit - start_jax_jit:.6f} 秒")

# 结果验证（考虑float32累积误差）
max_diff_conv = np.max(np.abs(output_np - output_jax_jit))
print(f"卷积结果最大绝对差值: {max_diff_conv:.6f}")
print(f"卷积结果近似程度 (atol=1e-3, rtol=1e-3): {np.allclose(output_np, output_jax_jit, atol=1e-3, rtol=1e-3)}")

# --- 结果可视化 ---
print("\n--- 输入输出图像可视化 ---")

fig, axes = plt.subplots(1, 2, figsize=(12, 6))

# 显示原始灰度图像
axes[0].imshow(image_np, cmap='gray', vmin=0, vmax=1)
axes[0].set_title('原始灰度图像 (输入)')
axes[0].axis('off')

# 显示模糊处理后的图像
axes[1].imshow(output_np, cmap='gray', vmin=0, vmax=1)
axes[1].set_title(f'模糊处理图像 (输出, 卷积核大小={kernel_size})')
axes[1].axis('off')

plt.tight_layout
plt.show

性能测试结果：

从以下路径加载图像: /mnt/d/images/taj_mahal.png
图像加载成功 (473x716)

--- 卷积性能基准测试 (图像: 473x716, 卷积核: 9x9) ---
可用JAX设备: [CudaDevice(id=0)]
NumPy (手动卷积)执行时间: 0.025815 秒
JAX (无JIT, 手动卷积)执行时间: 0.234791 秒
JAX (JIT, 手动卷积)首次运行时间（含编译）: 0.366345 秒
JAX (JIT, 手动卷积)第二次运行时间: 0.000238 秒
卷积结果最大绝对差值: 0.000000
卷积结果近似程度 (atol=1e-3, rtol=1e-3): True

再次验证了JAX JIT编译在第二次运行时实现约100倍的性能提升，展现了其在图像处理等计算密集型任务中的巨大潜力。