📚 课程目录第4章即时编译（jit）基础

4. 即时编译（jit）基础

JAX 的 jax.jit 是整套即时编译体系最常用的入口。它的作用不是“让 Python 变快”，而是把一段用 NumPy 风格写成的纯函数，在第一次遇到特定输入形状与静态配置时，追踪成一份 XLA 计算图，再交给后端编译成可在 CPU/GPU/TPU 上高效执行的机器码。理解 @jit 的关键，是理解“追踪—编译—缓存—复用”这条链路，以及哪些参数参与缓存键、哪些不参与。

一、@jit 的两种写法与基本约束

最直接的用法是装饰器形式，也可以对已有函数显式包装：

import jax
import jax.numpy as jnp

@jax.jit
def softmax(x):
    x = x - jnp.max(x, axis=-1, keepdims=True)
    e = jnp.exp(x)
    return e / jnp.sum(e, axis=-1, keepdims=True)

# 等价写法
fast_softmax = jax.jit(softmax)

x = jnp.array([[1.0, 2.0, 3.0],
               [0.5, 0.0, -1.0]])
y = softmax(x)          # 首次调用：追踪 + 编译
y2 = softmax(x * 2.0)   # 形状相同：直接复用已编译版本
print(y.shape, y.dtype) # (2, 3) float32

被 jit 包裹的函数应当尽量是“纯”的：输出只由输入决定，不依赖外部可变状态，不在内部做 Python 层面的副作用（如打印、写文件、修改全局列表）。原因是追踪阶段只执行一次 Python 代码，之后真正运行的是编译产物；若函数里含有依赖具体数值的 Python 分支，追踪时只会固化其中一条路径。例如 if x.sum() > 0: 在 jit 内会报错或行为不符合预期，应改用 jnp.where 表达。

二、编译缓存机制：缓存键由什么决定

每次调用 jit 函数，JAX 都会计算一个缓存键，主要包含：输入数组的 形状（shape）、数据类型（dtype）、静态参数值，以及被闭包捕获的常量结构。键相同则直接命中缓存，跳过追踪与编译；键不同则触发一次新的追踪编译。因此同一函数用 (32, 128) 和 (64, 128) 两种形状调用，会各自编译一份，这就是常见的“形状重编译”开销。

缓存默认按函数对象与键在进程内保存，可用 func.lower(...).compile() 预先编译，或用 jax.clear_caches() 在调试时清空。工程上应尽量让训练/推理的批大小、序列长度保持稳定，或把变化维度做成静态参数并接受有限几种组合，避免在循环里不断产生新形状。

三、静态参数与动态参数的区别

jit 默认把所有位置参数都当作动态参数：它们只影响运行时数值，不参与缓存键中的“值”，只贡献形状与 dtype。若某个参数的值会改变计算图结构（例如循环次数、分支选择、卷积核尺寸、是否启用某种算法），就应通过 static_argnums 或 static_argnames 声明为静态参数，让它的具体值进入缓存键。

@jax.jit
def poly(x, coeff):
    return coeff * x ** 2 + 1.0

# coeff 是动态参数：改值不重编译
poly(jnp.ones(4), 2.0)
poly(jnp.ones(4), 3.0)   # 复用同一编译版本

@jax.jit(static_argnames=("n_steps",))
def repeat_scale(x, n_steps):
    for _ in range(n_steps):   # 循环次数影响计算图
        x = x * 1.5
    return x

repeat_scale(jnp.ones(4), n_steps=3)  # 编译一次
repeat_scale(jnp.ones(4), n_steps=5)  # n_steps 变了 → 重新编译

判断原则可以概括为：值只作为张量参与运算的，用动态参数；值决定“图长什么样”的，用静态参数。静态参数必须是可哈希的（整数、字符串、元组、None 等），不能传数组。静态参数取值种类越多，缓存条目越多，编译次数也越多，因此不要把连续变化的浮点超参随意设为静态。

对比项	动态参数	静态参数
声明方式	默认行为	`static_argnums` / `static_argnames`
是否进入缓存键	只贡献 shape 与 dtype	具体值进入缓存键
改值是否重编译	否	是（值变化即新键）
典型对象	输入张量、权重、偏置	循环次数、模式开关、轴编号、算法选择
风险	无法用于控制 Python 分支	取值过多导致缓存膨胀

四、简单函数加速实战

下面用一个“逐元素复合运算 + 归约”的函数做对比示例。示例参数与耗时均为教学假设，实际数值随硬件与版本变化，重点观察调用方式与结果一致性。

import time
import jax
import jax.numpy as jnp

def heavy(x, w, b):
    h = jnp.tanh(x @ w + b)
    return jnp.mean(h * h, axis=-1)

fast = jax.jit(heavy)

key = jax.random.PRNGKey(0)
x = jax.random.normal(key, (1024, 256))
w = jax.random.normal(key, (256, 512)) * 0.05
b = jnp.zeros((512,))

# 预热：首次调用触发追踪编译
_ = fast(x, w, b).block_until_ready()

t0 = time.perf_counter()
for _ in range(50):
    out = fast(x, w, b)
out.block_until_ready()
t1 = time.perf_counter()
print("jit 平均耗时(ms):", (t1 - t0) / 50 * 1000)

# 与未编译版本对比
t2 = time.perf_counter()
for _ in range(5):
    ref = heavy(x, w, b)
ref.block_until_ready()
t3 = time.perf_counter()
print("普通平均耗时(ms):", (t3 - t2) / 5 * 1000)
print("最大绝对误差:", float(jnp.max(jnp.abs(out - ref))))

这段代码里有几个值得注意的实践点。第一，必须先预热：第一次调用包含追踪与编译时间，若把它算进计时会严重失真。第二，用 block_until_ready() 等待异步执行真正完成，否则计时只测到“派发任务”的时间。第三，编译版与普通版的结果应做数值比对，确认加速没有改变语义；由于浮点归约顺序可能不同，允许极小误差（如 1e-5 量级）。第四，若把 w 换成不同形状，会触发重编译，计时前应固定形状。

调试技巧：用 fast.lower(x, w, b).compile() 可提前完成编译并观察编译产物；用 jax.make_jaxpr(heavy)(x, w, b) 查看追踪出的中间表示，确认 Python 分支是否被正确固化。

常见误区：在 jit 函数内部用 print 观察数值、用 .item() 把数组转成 Python 标量、或依赖全局变量做累加，都会在追踪阶段出问题或只执行一次。需要条件逻辑时用 jnp.where，需要循环时用 jax.lax.fori_loop 或 jax.lax.scan。

把以上要点串起来：@jit 负责把纯函数编译成高效计算图，缓存键由形状、dtype 与静态参数值共同决定；动态参数改值不重编译，静态参数改值会重编译；实战中先预热、再计时、最后做数值一致性校验，才能得到可信的加速结论。