4、JAX函数式编程：纯函数与副作用、状态管理策略、参数更新模式

聊到JAX，就绕不开它的核心哲学——函数式编程。说实话，我第一次接触这个概念时，也觉得有点绕。毕竟我们做深度学习的，谁不是整天跟各种可变的状态、不断更新的参数打交道？但用久了你会发现，JAX这套设计，其实是在帮你规避很多隐形的坑。

4.1 纯函数与副作用：JAX的“洁癖”

先说说纯函数。纯函数有两个特点：同样的输入，永远得到同样的输出；不修改任何外部状态。说白了，它就是个“老实人”，不搞小动作。

举个例子，下面这个函数就不够“纯”：

import jax.numpy as jnp

# 不纯的函数 —— 依赖外部变量
scale = 2.0
def impure_fn(x):
    return x * scale  # scale 是外部变量，不是参数

# 纯函数 —— 所有依赖都通过参数传入
def pure_fn(x, scale):
    return x * scale

我在项目中遇到过这样的问题：刚开始用JAX时，习惯性地在函数里引用了一些全局配置变量。结果用vmap做向量化时，莫名其妙地报错。排查了半天，才发现是全局变量惹的祸。嗯，从那以后，我再也不敢在JAX函数里“偷懒”了。

4.2 状态管理策略：把状态“踢”出函数

既然函数不能有状态，那模型参数怎么办？总不能每次都重新初始化吧？

JAX的策略很简单：把状态当作参数传进去，再当作返回值传出来。你想想看，这不就是函数式编程的“三板斧”吗？

我习惯把模型参数封装成一个“参数字典”或“参数列表”，然后在训练循环中显式地传递它们：

import jax
import jax.numpy as jnp
from flax import linen as nn  # 虽然Flax不是必须的，但它很常用

# 定义一个简单的线性层
class SimpleLinear(nn.Module):
    features: int
    
    @nn.compact
    def __call__(self, x):
        # 这里的状态（参数）由Flax自动管理
        w = self.param('w', nn.initializers.normal(), (x.shape[-1], self.features))
        b = self.param('b', nn.initializers.zeros, (self.features,))
        return jnp.dot(x, w) + b

# 初始化参数 —— 状态被显式创建
model = SimpleLinear(features=10)
rng = jax.random.PRNGKey(0)
x = jnp.ones((5, 20))
params = model.init(rng, x)  # params 就是状态

# 前向传播 —— 状态作为参数传入
def forward(params, x):
    return model.apply(params, x)

output = forward(params, x)

你看，params被当作普通数据一样传来传去。这跟PyTorch那种“参数挂在模型对象上”的方式完全不同。刚开始可能不习惯，但用久了你会发现，这种显式的状态管理让代码的数据流变得非常清晰。

4.3 参数更新模式：函数式的梯度下降

参数更新是训练的核心。在JAX里，这个过程也是纯函数式的。我们来看看典型的更新模式：

# 定义损失函数 —— 纯函数
def loss_fn(params, x, y):
    pred = model.apply(params, x)
    return jnp.mean((pred - y) ** 2)

# 计算梯度 —— 也是纯函数
grad_fn = jax.grad(loss_fn)

# 更新参数 —— 返回新参数，不修改旧参数
def update(params, x, y, lr=0.01):
    grads = grad_fn(params, x, y)
    # 这里用到了“树结构”操作，JAX会自动遍历参数树
    new_params = jax.tree_map(lambda p, g: p - lr * g, params, grads)
    return new_params

# 训练循环 —— 显式地传递和更新状态
rng = jax.random.PRNGKey(42)
x_batch = jax.random.normal(rng, (100, 20))
y_batch = jax.random.normal(rng, (100, 10))

params = model.init(rng, x_batch)

for step in range(100):
    params = update(params, x_batch, y_batch)  # 每次返回新的params
    if step % 20 == 0:
        loss = loss_fn(params, x_batch, y_batch)
        print(f"Step {step}, Loss: {loss:.4f}")

这里有个关键点：jax.tree_map。它会递归地遍历参数树（比如嵌套的字典或列表），对每个叶子节点应用你指定的操作。说白了，就是帮你自动完成“参数-梯度”的配对更新，不用你手动写循环。

4.4 知识体系总览

为了让你更直观地理解JAX函数式编程的脉络，我画了一张图：

这张图把JAX函数式编程的三个核心概念串起来了。你会发现，它们其实都在强调同一件事：把状态当作数据来管理。纯函数保证了确定性，状态管理策略让参数可以显式传递，参数更新模式则展示了如何在这种范式下完成训练。

说实话，刚开始适应这套模式可能需要一点时间。但一旦你习惯了，就会爱上这种清晰的数据流。我在做分布式训练时，这种函数式设计让代码的并行化变得异常简单——因为所有状态都是显式的，不存在“哪个设备持有哪个参数”的混乱问题。

---