2、JAX基础快速回顾：JAX核心原语、函数式编程范式、不可变数组与随机数

好，咱们直接进入正题。如果你之前用过PyTorch或者TensorFlow，第一次接触JAX可能会有点懵。我刚开始也是这样——明明都是做深度学习的，怎么JAX连个nn.Module都没有？别急，这恰恰是它的魅力所在。

说白了，JAX就是一个「可微分的NumPy」。它把NumPy的数组操作、自动求导、GPU加速、以及自动并行化全部揉在了一起。你想想看，一个框架能同时搞定这些事情，那它的设计哲学一定很特别。

2.1 函数式编程范式：JAX的灵魂

JAX最让我着迷的一点，就是它彻底拥抱了函数式编程。什么意思呢？在JAX里，一切都是函数。模型是一个函数，损失是一个函数，优化器更新也是一个函数。

我在项目中遇到过这样一个坑：刚开始用JAX时，我习惯性地在函数内部修改全局变量，结果发现梯度计算完全不对。后来才明白，JAX要求你的函数必须是「纯函数」——没有副作用，同样的输入永远产生同样的输出。

import jax.numpy as jnp
from jax import grad, jit, vmap, pmap

# 纯函数示例：输入x，输出x的平方
def square(x):
    return x ** 2

# 错误示范：修改外部变量
# total = 0
# def bad_square(x):
#     global total
#     total += x ** 2  # 副作用！JAX会报错
#     return x ** 2

2.2 JAX四大核心原语

JAX有四个核心原语，掌握了它们，你就掌握了JAX的80%。我习惯把它们叫做「JAX四件套」：grad、jit、vmap、pmap。

原语	作用	类比
grad	自动求导	PyTorch的backward()
jit	即时编译加速	TensorFlow的@tf.function
vmap	自动向量化	手动写batch循环
pmap	自动并行化	多GPU数据并行

2.2.1 grad：自动求导

grad是JAX的自动求导工具。你给它一个函数，它返回这个函数的梯度函数。就这么简单。

from jax import grad

def f(x):
    return x ** 3 + 2 * x ** 2 + 1

# 获取f的梯度函数
grad_f = grad(f)

x = 3.0
print(f"f({x}) = {f(x)}")        # 输出：46.0
print(f"f'({x}) = {grad_f(x)}")  # 输出：39.0（3x² + 4x = 27 + 12）

你想想看，在PyTorch里你要写loss.backward()，然后从x.grad里取梯度。在JAX里，梯度就是一个函数，你直接调用它就行。我个人觉得这种方式更符合直觉——梯度不就是函数的导数吗？

2.2.2 jit：即时编译

jit是JAX的性能引擎。它把你的Python函数编译成高效的XLA代码，在GPU上跑得飞快。

from jax import jit
import time

def slow_function(x, y):
    for _ in range(1000):
        x = x * y + x
    return x

@jit
def fast_function(x, y):
    for _ in range(1000):
        x = x * y + x
    return x

x = jnp.ones((1000, 1000))
y = jnp.ones((1000, 1000))

# 第一次调用会编译，稍慢
start = time.time()
result = fast_function(x, y)
print(f"JIT第一次调用：{time.time() - start:.3f}秒")

# 后续调用飞快
start = time.time()
result = fast_function(x, y)
print(f"JIT第二次调用：{time.time() - start:.3f}秒")

2.2.3 vmap：自动向量化

vmap是我个人最喜欢的JAX原语。它能把一个处理单个样本的函数，自动变成处理batch的函数。说白了，就是帮你省掉写for循环的麻烦。

from jax import vmap

# 处理单个样本的函数
def process_single(x):
    return x ** 2 + 2 * x + 1

# 手动写batch处理
def process_batch_manual(batch):
    return jnp.array([process_single(x) for x in batch])

# 用vmap自动向量化
process_batch_vmap = vmap(process_single)

batch = jnp.array([1.0, 2.0, 3.0, 4.0])
print(process_batch_manual(batch))  # [4.0, 9.0, 16.0, 25.0]
print(process_batch_vmap(batch))    # 完全一样，但更快

为什么会更快？因为vmap不是真的在Python层面循环，而是把循环「推」到了XLA编译器里，生成高效的向量化代码。我在项目中用vmap处理过百万级别的数据点，速度比手动循环快了上百倍。

2.2.4 pmap：自动并行化

pmap是vmap的升级版。它把计算分布到多个设备（GPU/TPU）上并行执行。嗯，这里要注意，pmap需要你有多个设备才能发挥威力。

from jax import pmap

# 假设你有8个GPU
def train_step(params, batch):
    # 一个训练步骤
    return params - 0.01 * grad(loss)(params, batch)

# 自动在8个GPU上并行执行
parallel_train_step = pmap(train_step)

# 每个GPU处理自己的batch
params = init_params()
batches = split_data(data, num_devices=8)
params = parallel_train_step(params, batches)

2.3 不可变数组：JAX的「洁癖」

JAX的数组是不可变的。这意味着你不能像NumPy那样原地修改数组。我第一次遇到这个限制时，觉得JAX简直是在给自己找麻烦。但后来我明白了，这是为了函数式编程的纯洁性。

import numpy as np
import jax.numpy as jnp

# NumPy：可以原地修改
arr_np = np.array([1, 2, 3])
arr_np[0] = 100  # 没问题

# JAX：不能原地修改！
arr_jax = jnp.array([1, 2, 3])
# arr_jax[0] = 100  # 会报错！

# 正确做法：创建新数组
arr_jax_new = arr_jax.at[0].set(100)  # 返回新数组
print(arr_jax)      # [1, 2, 3]（原数组不变）
print(arr_jax_new)  # [100, 2, 3]（新数组）

你想想看，不可变数组的好处是什么？它让函数变得可预测。你传进去一个数组，函数绝对不会偷偷修改它。这对于调试和并行计算来说，简直是福音。

2.4 随机数：显式的随机状态

JAX的随机数生成器也遵循函数式编程的理念。它不像NumPy那样有一个全局的随机状态，而是要求你显式地传递和管理随机状态。

from jax import random

# 创建随机状态
key = random.PRNGKey(42)

# 生成随机数
key, subkey = random.split(key)  # 分裂出新的key
rand_vals = random.normal(subkey, shape=(3,))
print(rand_vals)

# 再次生成，需要新的key
key, subkey = random.split(key)
rand_vals_2 = random.normal(subkey, shape=(3,))
print(rand_vals_2)

2.5 知识体系总览

下面这张图是我自己整理的JAX核心知识体系，你可以把它当作一个快速参考地图。

这张图把JAX的核心概念串了起来。你从中心出发，先掌握四大原语，再理解三大基石，最后应用到实际场景中。我个人觉得，按照这个路径学习JAX，效率最高。

---