想象你拿到一本米其林三星餐厅的中央厨房菜谱：同样一斤牛肉、同样一勺生抽、同样 180 度火候，任何人都能做出同一条味道的糖醋里脊。菜谱不会因为厨师心情而改变味道，不会在做饭时偷偷把盐罐移到别处，更不会因为上次做过就影响下一次。这样的菜谱，就是干净、可靠、可被信任的。JAX 里的纯函数，就是这种“标准菜谱”。 可现实中的代码常常不像菜谱，而像一位随性的老厨师：边炒菜边尝味，顺手改一下墙上的盐罐，还时不时看眼手机决定放多少盐。结果同一道菜每次味道都对不上。JAX 作为一套为深度学习而生的函数式框架，把“纯函数”当作一切自动化的基石——只有函数无副作用、相同输入必得相同输出，JIT 编译、自动微分、并行计算才敢放心大胆地执行。 这篇文章会从一份菜谱出发，带你认识纯函数的定义、如何识别副作用、怎么把“不守规矩”的函数重构回纯函数，并理解它如何支撑 JAX 的三大杀手锏：JIT、grad 和 pmap。学完它能读懂 JAX 的文档，也能把日常训练代码写得更容易调试、加速和分布式扩展。

岗位关联

对 AI 从业者来说，掌握纯函数设计是驾驭 JAX 的门票：它决定你能否顺利使用 JIT 加速、自动微分和并行训练，也是写出可复现、可维护的深度学习代码的基本功。

目录
  1. 可视化演示
  2. 纯函数定义
  3. 副作用识别
  4. 函数式重构
  5. JIT编译前提
  6. 自动微分基础
  7. 并行计算保障
  8. 工程实践
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

动态演示JAX纯函数设计的核心机制，通过架构/数据/状态/性能四种模式展示关键概念与交互流程。

本节配有交互动效（进入视口后自动加载，约 3KB）。

JAX纯函数设计：纯函数的 纯函数的定义与核心特征 副作用与常见问题 核心机制可视化 JAX纯函数设计整体架构与关键组件

纯函数定义

中央厨房的菜谱，本质上是一张“输入—输出对照表”：给二两牛腩、一勺糖、半勺醋，出锅就是同一盘糖醋味。你不需要知道厨师张三是谁，也不需要关心外面是否下雨。菜谱的可信来自它只依赖原料，不借助任何天时地利。

JAX 希望每个函数都活成这样的菜谱。形式化地说，如果函数 f 的返回值只由它的参数决定，并且调用过程中不改变任何外部状态，那它就是一个纯函数。这也是 JAX 能放心对函数进行 JIT 编译、自动微分和并行加速的第一前提。

更严格地讲，纯函数要求“相同的输入一定得到相同的输出”。用数学语言表达：

就像同一位厨师用同一本菜谱、同一份原料炒十次，装盘必须一样；如果哪一次味道不同，多半是厨师偷尝了一口或偷偷改了火候。JAX 见到这种纯菜谱，才敢放心缓存、复制和重放。

Pure(f)  ⟺  ∀x1,x2 (x1=x2  ⇒  f(x1)=f(x2))\mathrm{Pure}(f) \iff \forall x_1, x_2 \,\big(x_1 = x_2 \;\Rightarrow\; f(x_1) = f(x_2)\big)
这是函数的确定性条件：只要输入对齐，输出就不会跑偏。
  1. 全称量词 \forall对任意两个可能的输入都做检查，不允许有例外。
    ∀x1,x2\forall x_1, x_2
  2. 输入相等判定如果两个参数的值完全相等，就进入下一步判断。
    x1=x2x_1 = x_2
  3. 输出相等结论那么两次调用得到的输出也必须相等；若不等，说明函数暗中依赖了隐藏状态。
    f(x1)=f(x2)f(x_1) = f(x_2)

副作用识别

厨师在炒菜时顺手尝了一口勺里的汁，或者在墙上记了一条“今天用了半瓶醋”，这些动作看似小，却会让同一道菜在别人复刻时对不上。菜谱之外的一切小动作，在函数式世界里统称副作用。

在 JAX 代码中，常见副作用有：打印日志、读写全局变量、原地修改数组、依赖随机数或当前时间。只要函数里出现其中任意一种，JAX 就无法确认它在不同调用之间是否保持一致。

用一个集合语言来约束：函数允许读取和写入的，只能是自己的参数。

判断一个 JAX 函数是否纯，就像检查一位厨师是否只在案板上操作：刀、锅、料都得在台面上；凡是伸进口袋摸手机看时间、往窗外喊话的，都不算守规矩。

Pure(f)  ⟺  reads(f)∪writes(f)⊆params(f)\mathrm{Pure}(f) \iff \mathrm{reads}(f) \cup \mathrm{writes}(f) \subseteq \mathrm{params}(f)
reads 和 writes 分别表示函数运行期间读取、写入的所有数据源，params 是显式参数集合。
  1. 读取集合 reads(f)函数运行期间读取的所有外部来源，比如全局数组、环境变量、系统时间。
    reads(f)\mathrm{reads}(f)
  2. 写入集合 writes(f)函数运行期间修改的所有外部来源，比如打印到 stdout、改动全局计数器。
    writes(f)\mathrm{writes}(f)
  3. 包含于 params(f)这些来源必须都被显式写在参数列表里；否则就是偷偷碰了外部盐罐。
    reads(f)∪writes(f)⊆params(f)\mathrm{reads}(f) \cup \mathrm{writes}(f) \subseteq \mathrm{params}(f)

函数式重构

发现厨师有“偷尝”和“改账本”的毛病后，该怎么改造成标准菜谱？不是没收他的勺子，而是把所有动作显式化：把尝的那勺糖也写进配料表，把墙上的账本变成一张待填写的回执，出锅时随菜一起交出。

函数式重构的核心是：把隐式依赖变成显式参数，把原地修改变成返回新对象。你不再“更新”一个状态，而是基于旧状态生成一个新状态，旧状态原地不动。

以状态转移为例，重构后的函数长这样：

在 JAX 里，这种“返回新对象”的模式天然贴合不可变数组。你每次更新参数，其实都是照着旧配方写出一份新配方；旧配方仍可留档、对比、回滚——这正是训练调试的宝贵能力。

s′=transition(s,a)且s 不被修改s' = \mathrm{transition}(s, a) \quad\text{且}\quad s \text{ 不被修改}
s 是旧状态，a 是输入动作，transition 返回新状态 s'，而旧状态 s 保持不变。
  1. 旧状态 s当前的状态对象，比如模型参数字典、环境变量列表。
    ss
  2. 输入动作 a这次调用携带的数据或决策，是计算出新状态的依据。
    aa
  3. 新状态 s'基于 s 和 a 生成的下一版本；原始的 s 不会少一个字段，也不会被改写。
    s′=transition(s,a)s' = \mathrm{transition}(s, a)

JIT编译前提

“预炒班”的故事：主厨拿到菜谱后，先在脑子里演一遍炒制流程，记住油温何时升高、何时下葱，之后每次只需按排练好的动作快速执行。这就像 JIT 的提前编译。

但如果菜谱里写着“看心情放盐”“根据窗外温度决定火候”，预演就毫无意义，因为每次执行都不一样。JAX 的 jit 要求函数可被静态追踪：类型、形状、控制流都能从输入确定，且不能依赖运行时的外部状态。

纯函数保证了 JIT 前后的行为等价，编译缓存才安全：

所以，在 JAX 里写纯函数，等于给编译器交了一份可预演菜谱。JIT 才会放心地把 Python 层的循环拍平成 XLA 计算图，下次直接按图快炒。

∀x:  fjit(x)=f(x)\forall x:\; f_{\mathrm{jit}}(x) = f(x)
不管编不编译，同一输入 x 必须得到同一结果；否则缓存可能给人错误答案。
  1. x 的任意性对所有可能的输入都要成立，而不是只在个别例子上相等。
    ∀x\forall x
  2. JIT 版本 f_jit经过 trace 和编译后的快速函数，本质是对同一菜谱的预演。
    fjit(x)f_{\mathrm{jit}}(x)
  3. 行为等价编译后的输出与原函数完全一致，这正是纯函数换来的信任状。
    fjit(x)=f(x)f_{\mathrm{jit}}(x) = f(x)

自动微分基础

如果一道菜最终味道偏甜，主厨想追责：是糖放多了，还是醋放少了？自动微分就是帮厨师按“配料变化→味道变化”的比例，把责任分摊到每种原料上。

可这份追责的前提是：配料到味道的因果关系不能被副作用搅乱。若厨师中途偷换了一瓶酱油，微分系统就分不清甜味来自配方的糖，还是来自误入的酱油。JAX 的 grad 基于计算图做反向传播，每一步都必须能忠实还原。

复合函数的链式法则正是自动微分的骨架：

纯函数让每个中间结果都有唯一确定的前因，反向传播才能像查账一样一追到底。若某一环出现随机读数或全局计数，账就对不上了——这就是为什么 JAX 的随机、日志都要放进参数里。

∂L∂x=∂L∂y⋅∂y∂x\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x}
损失 L 在中间量 y 上的导数，乘上 y 对 x 的导数，就把梯度传回了输入 x。
  1. 上层损失 L待优化的标量目标，比如均方误差、交叉熵。
    LL
  2. 中间结果 y前向传播中某个隐藏层或中间变量，它是连接 x 和 L 的环节。
    yy
  3. L 对 y 的偏导改变 y 会对损失产生多大影响。
    ∂L∂y\frac{\partial L}{\partial y}
  4. y 对 x 的偏导改变输入 x 会让 y 移动多少；两者相乘得到对 x 的梯度。
    ∂y∂x\frac{\partial y}{\partial x}

并行计算保障

承接前面所有铺垫，把菜谱分给十个厨师同时做：每人分到十桌的食材，各自按同一份配方操作，最后汇总成一百份一模一样的菜。并行系统最怕厨师之间偷偷传纸条、抢用同一个盐罐。

JAX 的 pmap / shard_map 会把数据切分到多设备，每个设备独立执行同一函数。只有函数没有写副作用、不碰全局状态，设备间才无需加锁，结果也不会因调度顺序而改变。

这可以写成逐元素不变性：

这份“各自独立成菜”的确定性，让大规模训练和推理可以放心横向扩展。以后你给 JAX 加数据并行，只要函数是纯的，基本不用改逻辑，纯粹增加厨师数量即可。

pmap(f)([x1,…,xn])=[f(x1),…,f(xn)]\mathrm{pmap}(f)([x_1,\dots,x_n]) = [f(x_1),\dots,f(x_n)]
把函数 f 映射到一批输入上，等价于对每个输入单独调用 f；并行只是换了分发方式，不改变任何单个结果。
  1. 输入列表n 个独立样本构成一个 batch。
    [x1,…,xn][x_1,\dots,x_n]
  2. pmap 调用JAX 将 f 复制到多个设备并对 batch 并行执行。
    pmap(f)([x1,…,xn])\mathrm{pmap}(f)([x_1,\dots,x_n])
  3. 逐元素结果整体结果等于把 f 分别套在每个 xi 上；设备间没有交叉副作用，所以可以安全拼接。
    [f(x1),…,f(xn)][f(x_1),\dots,f(x_n)]

工程实践

最后回到厨房管理：菜谱再纯，还要有统一的食材分类和出菜流程。JAX 的工程实践中，pytree 就像标准托盘——嵌套的字典、列表、元组都可以被统一遍历和映射，参数更新像给整桌菜统一换盘子。

训练中最常见的不可变更新，是用梯度去更新参数：不是把旧参数涂改掉，而是算出一份全新参数。这要求梯度计算与更新都遵循纯函数风格。

标准小批量梯度下降更新公式如下：

调试时，也把想看的中间量放进函数返回值里，而不是 print 到控制台；所有输入输出都通过参数与返回值传递。这样整个训练过程像一张张可留档的菜谱，哪里不对就回看哪一版——这就是纯函数设计给工程留下的最大红利。

θt+1=θt−η⋅∇θL(θt)\theta_{t+1} = \theta_t - \eta \cdot \nabla_\theta L(\theta_t)
当前参数减去学习率乘以梯度，得到下一轮参数；整棵 pytree 上的每个叶子都会按同一规则更新。
  1. 当前参数 θt第 t 步模型参数，可以是嵌套的 pytree。
    θt\theta_t
  2. 学习率 η控制每步更新幅度，就像决定每次放多少盐的勺量。
    η\eta
  3. 梯度方向JAX 的 grad 返回损失对参数的梯度，指示损失上升最快的方向。
    ∇θL(θt)\nabla_\theta L(\theta_t)
  4. 新参数 θt+1沿着梯度反方向走一小步，得到下一轮的状态。
    θt+1=θt−η⋅∇θL(θt)\theta_{t+1} = \theta_t - \eta \cdot \nabla_\theta L(\theta_t)

AI 应用场景

大规模模型训练中的损失函数

阶段：模型训练

解决的问题：传统命令式训练代码里，损失函数容易悄悄依赖全局 step、外部参数副本或日志打印，导致多次运行结果不一致；而 JAX 要求损失函数是纯函数，这样 grad 和 jit 才能稳定工作。

具体案例：训练一个图像分类模型时，把损失函数写成 loss(params, batch)。params 和 batch 都从参数传入，函数内部只用 jnp 计算交叉熵。这样 jax.jit(loss) 可以缓存编译结果，训练速度提升 30% 以上。

分布式数据并行推理

阶段：分布式推理或训练

解决的问题：多卡并行时，如果每个设备共享一个全局随机状态或公共缓存，会出现结果不确定、通信竞争等问题。纯函数让每个设备只依赖自己的分片输入。

具体案例：在 8 张 GPU 上用 pmap 并行推理一批文本。每个设备运行同一个纯函数 forward(params, tokens)，设备间不读写共享变量，最后把输出按 batch 维度拼接，保证结果与单卡完全一致。

科学仿真中的可复现实验

阶段：实验管理

解决的问题：科研实验需要严格复现，任何读时间、写全局计数器的代码都会让结果“随缘”。纯函数配合显式随机 key，让随机性也能被复现。

具体案例：强化学习仿真中，把环境状态、动作和随机 key 都作为参数传入 step_fn(state, action, key)。每次返回新状态和新 key，日志统一记录状态快照，实验可以逐帧重放。

Python 演示

python30 行
import jaximport jax.numpy as jnp global_counter = 0 def pure_loss(params, x, y):    pred = jnp.dot(x, params)    return jnp.mean((pred - y) ** 2) def impure_loss(params, x, y):    global global_counter    global_counter += 1    pred = jnp.dot(x, params) + global_counter * 0.01    return jnp.mean((pred - y) ** 2) key = jax.random.PRNGKey(42)x = jax.random.normal(key, (8, 3))y = jax.random.normal(key, (8,))params = jax.random.normal(jax.random.PRNGKey(1), (3,)) print("pure call #1:", pure_loss(params, x, y))print("pure call #2:", pure_loss(params, x, y))print("impure call #1:", impure_loss(params, x, y))print("impure call #2:", impure_loss(params, x, y)) grads = jax.grad(pure_loss)(params, x, y)print("pure grads:", grads) fast_loss = jax.jit(pure_loss)print("jit loss:", fast_loss(params, x, y))
目的：用最小例子展示纯函数与带副作用函数在 JAX 中的表现差异，并演示 grad 与 jit 如何依赖纯函数。
  1. def pure_loss(params, x, y):纯函数版本，所有输入来自参数，没有外部状态。
  2. global global_counter在函数体内修改外部变量，是 impure_loss 产生副作用的根源。
  3. global_counter += 1每次调用都会改变外部计数器，导致相同输入下输出不同。
  4. jax.grad(pure_loss)(params, x, y)JAX 自动求导，纯函数可安全地重建计算图。
  5. jax.jit(pure_loss)通过 JIT 编译缓存，纯函数保证编译前后行为一致。

常见误区

误区print 只是打印点文字，不影响数值计算，所以就算函数里有 print 也没什么大不了的。

正解print 是标准 I/O 副作用。在 jit 或 grad 中，print 可能只在 tracing 时执行一次，也可能被编译器完全忽略，导致你看到的调试信息和实际运行过程对不上。纯函数要求所有对外沟通都走返回值，调试时可以把中间量打包返回，或用 jax.debug.print 这种专为 JAX 设计的打印方式。

为何常见：很多初学者从命令式 Python 转来，习惯用 print 观察中间量，忽略了 JAX 的 trace 机制。

误区函数式更新要不断生成新数组，比原地修改慢多了，纯函数只适合教学不适合生产。

正解JAX 的数组本身不可变，原地修改并不被鼓励。XLA 编译器在 jit 下会做缓冲区复用和内存分析，纯函数返回的新数组常常可以复用在旧缓冲区上，并不见得产生额外拷贝。真正的性能瓶颈往往来自不纯的函数导致无法编译，而不是“返回新对象”本身。

为何常见：从 NumPy 的 in-place 操作习惯迁移过来后，容易误判内存分配的开销，忽略编译优化带来的收益。

误区我的函数没有赋值给任何外部变量，只读了一些全局配置，应该还算纯函数吧。

正解纯函数不仅要求“不修改”，还要求“不读取”函数作用域之外的动态状态。读取全局配置、依赖当前时间、调用随机数生成器、执行 I/O 都会让相同参数可能产生不同输出。正确做法是把这些依赖全部作为参数传入，例如 JAX 的随机 key。

为何常见：人们容易把“副作用”简单理解为“修改外部状态”，忽略了读取外部可变状态同样会破坏确定性。

误区深度学习离不开 dropout、初始化、数据采样这些随机操作，纯函数根本做不到。

正解随机性本身没有问题，问题在于随机来源是否可控。JAX 用显式的 PRNGKey 作为参数传递给随机函数，相同的 key 产生相同的随机序列；这样随机操作也是纯函数。于是我们可以让实验可复现，也能安全地把随机过程放进 jit 和 grad 的计算图中。

为何常见：混淆了“算法的随机性”与“实现上的不可控副作用”，以为随机就一定等于非确定性。

面试题

Q1. 关于 JAX 纯函数，下列说法正确的是？

单选题 · easy

  • A.相同输入可能得到不同输出
  • B.函数可以修改全局计数器
  • C.函数只依赖参数产生返回值← 正确答案
  • D.函数必须返回字典类型

答案：C。纯函数要求返回值只由参数决定，并且不修改外部状态。A 和 B 都违背了确定性，D 是错误的额外限制，只有 C 符合定义。

Q2. 下列哪一项属于 JAX 中的副作用？

单选题 · easy

  • A.返回一个全新的数组
  • B.调用 jnp.sin 计算正弦
  • C.打印一条日志到终端← 正确答案
  • D.使用参数做矩阵乘法

答案：C。返回新数组、调用纯算子、用参数计算都是合法的函数式行为；打印日志属于 I/O，会带来外部可观察效果，因此是副作用。

Q3. 在 JAX 中，随机函数需要显式传入 PRNGKey 是为了什么？

单选题 · easy

  • A.让随机性变得更加强烈
  • B.使随机函数变成纯函数← 正确答案
  • C.让训练过程明显变慢
  • D.避免使用任何 NumPy 函数

答案：B。把随机源作为 key 参数传入，相同 key 得到相同随机序列，函数不再依赖隐式的全局随机状态，因此满足纯函数定义。A、C、D 都不是设计目的。

Q4. JIT 编译能够安全缓存函数结果的前提是什么？

单选题 · mid

  • A.函数内部代码足够短小
  • B.被编译的函数是纯函数← 正确答案
  • C.函数使用全局配置变量
  • D.函数返回的是一个浮点数

答案：B。JIT 需要函数可重放、可静态追踪。只有纯函数能保证同一输入在任何时候调用都得到同一输出，缓存才有效；代码长度、返回类型都不是关键。

Q5. 自动微分中链式法则的主要作用是什么？

单选题 · mid

  • A.把损失函数的值降为零
  • B.把梯度从输出逐层传回输入← 正确答案
  • C.让反向传播不再需要计算图
  • D.跳过所有中间变量直接求导

答案：B。链式法则把损失对中间量的梯度与中间量对输入的梯度相乘，从而让梯度逐层回传。它正是反向传播实现的基础，不是用来消灭损失或跳过中间量。

Q6. 若函数 f 读写了全局计数器，要重构为纯函数，最合适的做法是？

单选题 · hard

  • A.把计数器改成全局列表保存
  • B.将计数器作为参数传入并返回新值← 正确答案
  • C.删除计数器变量并重写逻辑
  • D.用打印日志代替计数器累加

答案：B。函数式重构的关键是显式化：counter 作为参数进入函数，函数返回 （结果，新counter）。这样函数不碰外部状态，调用方自己维护状态流转。A、D 仍然有副作用，C 可能丢失功能。

Q7. 在 JAX 中使用 pmap 并行执行一个纯函数，下列哪项正确？

单选题 · mid

  • A.不同设备上的结果会偶然不同
  • B.所有设备必须共享一个全局变量
  • C.各设备独立执行，结果可预测拼接← 正确答案
  • D.pmap 只能在单张显卡上运行

答案：C。纯函数让每个设备只依赖自己的分片输入，互不干扰，因此结果确定且可以拼接。全局共享状态反而会引入竞争条件，pmap 本来就是多设备并行 API。

Q8. 下列哪些写法会让 JAX 函数丧失纯函数性？

多选题 · mid

  • A.打印输入数组的形状← 正确答案
  • B.修改传入的字典参数← 正确答案
  • C.调用 jnp.dot 计算内积
  • D.读取 CPU 当前时间← 正确答案

答案：ABD。打印日志是 I/O，修改传入参数是写外部状态，读取当前时间让结果依赖隐藏环境；这三者都属于副作用。jnp.dot 是纯算子，不会破坏纯函数性。

Q9. 纯函数为 JAX 带来的核心能力包括哪些？

多选题 · hard

  • A.JIT 编译可以安全缓存结果← 正确答案
  • B.自动微分能够忠实追溯梯度← 正确答案
  • C.并行计算保持结果确定性← 正确答案
  • D.自动生成训练用的数据集

答案：ABC。纯函数让计算可重放，因此 JIT 缓存、梯度追溯和并行拼接都变得可靠。但纯函数本身不会自动生成数据，这属于数据处理流程，不是纯函数带来的核心能力。

Q10. 函数式重构中，处理可变状态的可取做法有哪些？

多选题 · mid

  • A.把状态作为参数显式传入← 正确答案
  • B.返回新状态而不是原地修改← 正确答案
  • C.使用全局缓存保存中间状态
  • D.把随机 key 作为参数传递← 正确答案

答案：ABD。显式传参、返回新状态、显式传递随机 key 都是常见的函数式重构手段。全局缓存是隐藏外部状态，会重新引入副作用，因此不可取。