📚 课程目录第3章纯函数与副作用管理

3. 纯函数与副作用管理

JAX 的 jit、vmap、grad 等变换之所以能自由组合，根源在于它假设被变换的函数是纯函数：相同输入必然得到相同输出，且函数内部不修改外部可见状态。理解这条约束，是写出可编译、可向量化、可求导代码的前提。本节从“为什么必须纯”讲起，落到随机数这一最典型的“伪副作用”场景，再给出状态传递与副作用规避的工程做法。

一、JAX 为什么要求纯函数

JAX 的变换不是运行时解释执行，而是对函数做追踪（tracing）：用抽象的 Tracer 对象代替真实数组走一遍计算图，记录下算子序列，再交给 XLA 编译。追踪只发生一次，之后同一编译产物会被反复调用。如果函数依赖外部可变状态，问题立刻出现：

结果不可复现：追踪时读到的外部变量值会被“固化”进计算图，后续外部变量改变，编译产物不会跟着变。
副作用被丢弃或重复：print、写文件、修改全局列表等操作在追踪阶段可能只执行一次，或在 vmap 下被批量复制，语义完全失控。
无法求导与并行：grad 需要函数的数学表达，vmap 需要把标量运算提升为批量运算，任何隐藏的状态读写都会破坏这种代数结构。

因此 JAX 的实践原则是：把“状态”从函数内部搬到函数参数与返回值上，让函数只做“输入 → 输出”的映射。凡是需要变化的东西，都显式地作为参数传入、作为返回值传出。

二、随机数：最典型的“伪副作用”

NumPy 的 np.random.randn() 依赖一个全局随机状态，每次调用都推进它，因此同一行代码两次调用结果不同——这是典型的非纯行为。JAX 彻底改掉了这个模型：随机数由显式的密钥（PRNGKey）决定，密钥是输入，随机数是输出，同一个密钥永远产生同一串随机数。

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)          # 由整数种子生成密钥
a = jax.random.normal(key, (3,))
b = jax.random.normal(key, (3,))
print(jnp.allclose(a, b))            # True：同 key 同结果

# 需要“新”的随机数时，必须显式分裂密钥
key, subkey = jax.random.split(key)
c = jax.random.normal(subkey, (3,))
print(jnp.allclose(a, c))            # False：不同 key 不同结果

关键点有三个。第一，绝不要复用同一个 key 去生成两组本应独立的随机数，否则它们会完全相同，在初始化多层网络参数时会造成隐蔽的对称性问题。第二，用 split 派生新密钥，而不是对 key 做加法或自增，因为 split 保证派生出的子密钥在统计上互不相关。第三，key 必须作为参数在函数间传递，不能藏在闭包或全局变量里。

常见错误：在 jit 函数内部写 key = jax.random.PRNGKey(0)。这会让每次调用都从同一个种子开始，随机性完全丧失。正确做法是把 key 作为函数入参传入。

三、状态传递：把“变化”变成“数据流”

训练循环里的参数更新、优化器动量、Dropout 掩码、批次的随机打乱，都是需要跨步骤演进的状态。JAX 的解法是函数式状态传递：状态是普通的数据结构（通常是 pytree，如字典或具名元组），每一步由函数接收旧状态、返回新状态。

def train_step(params, opt_state, key, batch):
    key, drop_key = jax.random.split(key)
    loss, grads = jax.value_and_grad(loss_fn)(params, batch, drop_key)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, key, loss

# 循环中显式接收并回传全部状态
for batch in batches:
    params, opt_state, key, loss = train_step(params, opt_state, key, batch)

这段代码里没有任何隐藏状态：params、opt_state、key 都是显式入参与返回值，train_step 因此可以被 jit 编译，也可以在 vmap 下对多组参数并行。注意 key 在函数内先 split 再使用，并把推进后的 key 返回，这样下一次调用拿到的是新密钥，随机流才能持续演进。

当状态很多时，建议把它们打包成一个 pytree 统一传递，避免参数列表膨胀：

state = {"params": params, "opt_state": opt_state, "key": key, "step": 0}

def step(state, batch):
    key, sub = jax.random.split(state["key"])
    ...
    return {**state, "params": new_params, "key": key, "step": state["step"] + 1}

四、副作用规避策略

真实工程中总有些操作天生带副作用：日志、指标记录、文件写入、原地修改数组。JAX 提供了分层的手段来处理它们，选择哪种取决于副作用是否影响计算结果。

场景	推荐做法	说明
调试打印	`jax.debug.print`	在 `jit` 内可用，按运行时真实值打印，替代普通 `print`
指标收集	作为返回值传出	把 loss、准确率等放进返回值，在循环外统一记录
数组原地更新	`arr.at[i].set(v)`	函数式更新，返回新数组，语义等价且可追踪
必须的外部写入	`jax.experimental.io_callback`	显式声明为宿主回调，明确其非纯性质
随机性	PRNGKey 显式传递	把随机性建模为纯函数，见上文

其中 .at[] 语法值得单独强调。习惯 NumPy 的人会写 arr[i] = v，但在 JAX 中这会报错或行为异常，因为 JAX 数组不可变。arr.at[i].set(v) 返回一个新数组，旧数组不变，这正是纯函数语义。类似地还有 .at[i].add(v)、.at[idx].get() 等，覆盖了绝大多数原地操作需求。

判断一个函数能否安全交给 jit 的快速自检：把函数体里所有对外部变量的读写列出来，逐个问“它是否作为参数传入或作为返回值传出”。如果答案都是“是”，函数基本就是纯的；如果有“否”，就需要按上表改造。

五、一个完整的对照示例

下面把“非纯写法”和“纯函数写法”放在一起对比，差异一目了然：

# 反例：依赖全局状态，无法 jit
rng_state = 0
def bad_sample(n):
    global rng_state
    rng_state += 1
    return jax.random.normal(jax.random.PRNGKey(rng_state), (n,))

# 正例：key 显式传入传出，可 jit、可 vmap
@jax.jit
def good_sample(key, n):
    return jax.random.normal(key, (n,))

key = jax.random.PRNGKey(42)
keys = jax.random.split(key, 4)
samples = jax.vmap(good_sample, in_axes=(0, None))(keys, 3)  # 形状 (4, 3)

正例中 good_sample 是纯函数，因此既能被 jit 编译，又能被 vmap 沿 key 维度批量映射，一次得到 4 组独立样本。反例中的全局变量在追踪时只会被读取一次，后续调用完全失去随机性——这正是 JAX 强制纯函数所要避免的陷阱。