你有没有注意过：有些厨师做菜全凭感觉，围裙口袋里塞着各种私藏调料，上次多放了盐，这次就可能少放糖——你永远没法预测他端出来的是什么味道。而另一些餐厅，每道菜的配方都写成一张“菜谱卡”：食材、火候、调料用量，全部写得明明白白。只要食材一模一样，任何厨师照着做，出来的味道也一模一样。这张菜谱卡，就是计算机世界里的“纯函数”。 这篇文章要讲的 JAX 纯函数范式，正是把深度学习的模型代码改写成一张张“菜谱卡”：模型参数不再是藏在对象口袋里的秘密，而是像食材一样摆在桌面上，显式地传给函数；训练过程也不是把旧菜谱涂改掉，而是拿着旧菜谱生成一张新菜谱。你会看到，这种“把事情摊开做”的风格，让代码变得可复现、可并行、可组合，也让 JAX 的自动微分、即时编译等强大能力有了施展的舞台。

岗位关联

掌握纯函数范式，意味着你能写出可复现、易并行、好调试的深度学习代码，同时吃透 JAX 的设计哲学，为迁移到函数式深度学习、构建大规模训练系统打下坚实基础。

目录
  1. 可视化演示
  2. 纯函数本质
  3. 显式状态传递
  4. 不可变更新
  5. 模块化设计
  6. PyTorch对比
  7. 工程优势
  8. 自动微分
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

展示标记-清除垃圾回收算法的执行过程：从GC根出发标记可达对象，然后清除不可达对象释放内存。

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

初始状态 对象图中包含GC根（R1/R2）和8个堆对象 部分对象已无引用路径，等待回收 GC Roots = Stack + Globals + Registers

纯函数本质

回到那家餐厅。有张菜谱卡，写的是“西红柿炒蛋：番茄两个、鸡蛋三枚、盐三克，中火三分钟”。请问：换个厨师，食材不变，味道会变吗？不会。你可以把这张卡复印一百份，寄到一百家分店，只要买到的番茄鸡蛋一样，出锅的菜就一模一样。在编程世界里，这种“只靠输入就能完全决定输出，且绝不碰外界一根手指”的函数，就叫纯函数。

深度学习的代码，在大多数框架里并不纯粹。模型训练时可能悄悄改了全局日志、依赖了系统时间、或者在某次前向传播里顺手更新了缓冲区。而 JAX 把“纯函数”当作第一公民：每个算子、每个模型，最好都写成“给什么返回什么”的形态。用数学语言说，一个纯函数就是一条从输入集合到输出集合的确定映射。

纯函数有两条铁律：第一，引用透明——同一个表达式在任何地方求值，结果都一样；第二，无副作用——不修改全局变量、不写文件、不依赖外部状态。对深度学习而言，这意味着：同一份参数 $\theta$ 和同一样本 $x$，loss 永远一样。这正是复现实验、定位 bug、并行计算的根基。JAX 的 jit 变换会把函数重新执行很多次，如果函数里藏了副作用（比如打印日志），日志就会被重复输出，这种“不干净”的写法在函数式世界观里是头号大忌。

所以，纯函数就是一张“老实巴交”的菜谱卡：它不偷吃食材，不弄脏厨房，更不会因为今天厨师心情好就多放一勺糖。学会辨认和书写纯函数，是你走进 JAX 世界的第一步——接下来，我们就要把藏着掖着的模型状态，一张一张地摆到餐桌上来。

y=f(x)\displaystyle y = f(x)
一个纯函数就像一张诚实的菜谱卡：输入食材 $x$，经过固定的做法 $f$，必然得到唯一的菜品 $y$，中间不偷吃、不乱改、不依赖厨房环境。
  1. x（输入食材）函数唯一的输入来源。在神经网络中，它可以是样本特征、模型参数等一切外部数据。
    xx
  2. f（映射规则）从输入到输出的确定变换。对应前向传播函数，内部只做计算，不看、不改任何外界的东西。
    f(⋅)f(\cdot)
  3. y（确定输出）由 x 唯一决定的输出。同一个 x 在任何时刻、任何设备上调用，得到的 y 都完全相同，这就是可复现性。
    y=f(x)y = f(x)

显式状态传递

传统的面向对象厨房里，厨师会把私藏调料塞在围裙口袋里。你喊一声“做道菜”，他翻翻口袋自己找料，你根本不知道他用了什么。JAX 的哲学恰恰相反：所有食材必须全部摆在桌面上，一眼看尽。模型参数 $\theta$ 不再是对象的内部属性，而是作为一个普通参数，明明白白地写进函数签名里。

在 PyTorch 中，你写 model(x) 就能完成前向传播，因为 model 对象内部自己保存着权重。但在 JAX 中，标准的写法是 f(params, x)：第一份资料是模型参数，第二份才是输入数据。没有隐藏状态，没有“函数内部自己知道的东西”——所有变量都来自参数，所有结果都流向返回值。

那么多个参数怎么传递？JAX 提供了 Pytree 结构，它像一个多格子的收纳盒：把神经网络的权重、偏置、甚至优化器的动量状态，全部塞进同一个盒子，作为单个参数传入函数。函数内部需要用哪些格子，就打开哪个格子取。收纳盒让“显式传递”不再繁琐，反而比一堆散落的全局变量更整洁、更有条理。

把食材从口袋里掏出来摆在桌上，一开始可能不习惯，但好处立刻显现：任何人看到 f(params, x) 这个签名，就知道这个函数需要什么；任何人看到返回值，就知道它产出了什么。状态不再神出鬼没，代码变得像菜谱卡一样容易被“读懂和复现”。

y=f(θ,x)\displaystyle y = f(\theta, x)
菜谱卡写得很清楚：食材 $\theta$ 和 $x$ 一起下锅，成品 $y$ 出锅。没有一口“看不见的锅”，也没有任何“备用调料”。
  1. $\theta$（显式参数）模型权重作为第一个参数显式传入，不再隐藏在对象内部。
    θ\theta
  2. x（样本输入）当前批次的数据，与参数并列作为函数输入。
    xx
  3. f（纯函数）只依赖 $\theta$ 和 $x$，不读取任何外部状态，也不修改任何外部变量。
    f(⋅)f(\cdot)
  4. y（输出结果）预测值或损失值，通过返回值流出函数，而不是写进某个全局位置。
    y=f(θ,x)y = f(\theta, x)

不可变更新

假设餐厅老板想改良菜谱，他会怎么做？把旧菜谱卡涂掉重写吗？不会——他更可能先复印一张，在新副本上改：万一新配方失败，旧配方还完好无损。JAX 的训练循环正是这种“复印改稿”的思路：参数一旦创建，就永远不被修改；每次更新，都是基于旧参数生成一份全新的参数结构。

这就是“函数式更新”。传统写法里，optimizer.step() 会悄悄地把 model.weight 原地修改掉；而 JAX 中，你写 params = update(params, grads)，右边用旧参数计算出新参数，左边的变量名虽然还是 params，但指向的是全新的对象。旧参数对象依然存在于内存中，随时可以拿出来回溯、对比、甚至继续训练。

随机数也要这样管理。在 Python 里，random.seed(0) 之后调用 random.random()，生成器内部状态在悄悄变化——这不是纯函数。JAX 的做法是把随机数“钥匙” key 当作显式参数传递，每次用完就用 jax.random.split(key) 分裂出新钥匙，旧钥匙永远保持原样。这样，每一次随机采样都变成了可复现的纯函数调用。

“旧卡永远不改”听起来像浪费内存，但正是这种不可变性，让训练过程像一本拍立得相册：每一轮的参数都是一张永久保存的照片。你可以翻回第 1200 轮的照片重新推理，也可以把同一张照片发给一百家分店同时试菜——这就是可复现与可并行的秘密。

θ′=θ−η g\displaystyle \theta' = \theta - \eta\, g
一张菜谱卡（$\theta$），根据营养分析报告（梯度 $g$）和学习步长（$\eta$），复印出一张新卡（$\theta'$）。旧卡不被涂抹，新卡完全独立。
  1. $\theta$（旧参数）当前训练轮次的参数，保持不可变，作为更新过程的输入。
    θ\theta
  2. g（梯度）损失对参数的偏导，指明参数应该往哪个方向调整。
    gg
  3. $\eta$（学习率）每步更新的步长标量，控制参数移动的幅度。
    η\eta
  4. $\theta'$（新参数）更新后生成的全新参数结构，旧参数对象没有被修改。
    θ′=θ−η g\theta' = \theta - \eta\, g

模块化设计

一张完整的宴席菜单，从来不是一句话写成的，而是由几十张小菜谱卡组合而成：前菜、主菜、甜点、汤羹，各司其职。JAX 的模块化设计也是一样：复杂的神经网络，由一个个纯函数小块组合而成，每个小块只负责一件事。

组合的方式就是“函数复合”。如果 f 是特征提取层，g 是分类层，那么整个模型就是先经过 f 再经过 g，写成数学记号就是 $g \circ f$。在代码里，你甚至不需要特殊的容器类，直接写 h = lambda x: g(f(x)) 就完成了组合。每个小函数都可以独立测试、独立复用，就像单独检验一张菜谱卡的正确性。

JAX 还大量使用高阶函数——把“变换”当作参数传递的函数。比如 jax.grad 接收一个函数，返回它的梯度函数；jax.jit 接收一个函数，返回编译加速后的版本。这种“函数吃函数”的玩法，让能力的组合变得极其灵活：你可以把纯函数理解成一个积木，高阶变换就是积木之间的通用接口，任意积木都能咬合。

积木式的设计让团队协作变得顺畅：每个人只负责打磨几张“小菜谱卡”，然后用统一的接口把它们拼起来。哪张卡出了问题，单独测试那一张就能定位；哪张卡想升级，换一张新卡接上去即可，不影响厨房里的其他任何东西。

h=g∘f\displaystyle h = g \circ f
把小菜谱卡 $f$ 和 $g$ 首尾相接，就得到一张新菜谱卡 $h$。拼接过程不会修改任何一张原卡，新的组合卡依然是一张纯粹的菜谱。
  1. f（前段工序）第一个小纯函数，接收原始输入并产生中间结果。
    ff
  2. g（后段工序）第二个小纯函数，接收中间结果并产生最终输出。
    gg
  3. $\circ$（复合操作）函数复合的数学符号，表示“先做 f，再做 g”。
    ∘\circ
  4. h（组合函数）复合后得到的新函数，整体依然是一个纯函数，可以继续参与复合。
    h=g∘fh = g \circ f

PyTorch对比

到这一步，你可能想问：那 PyTorch 里的模型该怎么对应？在 PyTorch 里，模型是一位“围裙厨师”：net = nn.Linear(10, 1)，权重存在 net.weight 这个围裙口袋里，你只管喊 net(x)，他翻翻口袋自己找参数。而 JAX 的菜谱卡风格是：把 $\theta$ 摆上餐桌，写清楚 f(params, x)。

两种风格都能做出同一道菜，但状态存放的位置决定了代码的属性。nn.Module 将参数作为对象属性保存在 self 里，调用 model(x) 时隐式读取；JAX 则要求所有可变状态全部显式出现在函数签名中。对比来看：

从 PyTorch 迁移到 JAX，最需要改变的不是语法，而是心智模型：从“我是谁、我有什么属性”切换成“数据怎么流、变换怎么组合”。实操上的迁移路径很简单：第一步，把你模型里所有 self.xxx 列成一张状态清单；第二步，把这些状态打包成 Pytree，作为函数的第一个参数；第三步，把每个方法改写成“参数进、结果出”的纯函数。三步走完，你的代码就脱下了围裙，换上了菜谱卡。

yOO=model(x)vsyFP=f(θ,x)\displaystyle y_{\text{OO}} = \text{model}(x) \quad \text{vs} \quad y_{\text{FP}} = f(\theta, x)
左边是在喊一位带围裙的厨师：你只管下单，参数他自己找；右边是把菜谱卡和食材一起递出去：所有输入都在明面上。
  1. model(x)（对象式调用）PyTorch 中模型是带状态的对象，调用时偷偷读取内部 self 里的权重。
    model(x)\text{model}(x)
  2. f(θ, x)（函数式调用）JAX 中参数 θ 与输入 x 并列，作为显式参数传入函数。
    f(θ,x)f(\theta, x)
  3. 状态位置对比两者算出的 y 可能完全等价，但依赖关系一个藏在对象里，一个摆在明面上。
    yOO  vs  yFPy_{\text{OO}} \;\text{vs}\; y_{\text{FP}}

工程优势

连锁餐厅的老板最头疼什么？一是每家分店口味不稳定，二是新菜推广太慢。菜谱卡模式完美解决了这两个问题：口味不稳定？同一张卡加同一批食材，味道必然一样；推广太慢？菜谱卡复印一千份，一千家分店同时开火。深度学习的工程实践，面临的正是同样的两个难题：训练不可复现、并行扩展困难。

纯函数范式让“可复现训练”变成一件自然的事。参数不可变、随机数显式管理、函数无副作用——训练过程中的每一个中间状态都被记录、可回溯。上周跑出来的 loss 曲线，这周重新跑一遍，几乎逐位一致。调试时，你只需要保存当时的参数快照和输入数据，就能像回放录像一样还原现场，不需要猜测“那个对象当时的状态是什么”。

并行方面，纯函数天然安全：多个设备、多个线程可以同时调用同一个函数，因为它们互不干扰、不共享可写状态。批量训练时，每个样本的损失计算都是独立的纯函数调用，把它们的结果平均起来就是整个 batch 的损失：

除此之外，无副作用还意味着“缓存友好”：同样的输入，输出必然相同，所以计算结果可以安全地被缓存复用；分布式训练时，梯度同步不需要加锁，因为不存在共享的可变状态。一套菜谱卡，既保证了口味统一，又让千百家分店可以同时开工——这正是工业级深度学习最需要的两个特质。

L=1N∑i=1Nℓ(f(θ,xi),yi)\displaystyle \mathcal{L} = \frac{1}{N}\sum_{i=1}^{N}\ell\big(f(\theta, x_i), y_i\big)
一桌客人点的菜各不相同，但每张菜谱卡都是独立的：厨师们可以同时开火，最后把每道菜（每个样本的损失）平均一下，就是一桌宴席的整体评价。
  1. N（样本数量）当前批次中的样本总数，对应一桌客人的数量。
    NN
  2. $\ell$（单样本损失）衡量单个预测与真实标签差距的损失函数。
    ℓ(⋅)\ell(\cdot)
  3. $\sum$（求和）把每个样本的损失加起来，代表整桌菜的总分。
    ∑i=1N\sum_{i=1}^{N}
  4. $\frac{1}{N}$（平均）除以样本数得到平均损失，作为训练优化的目标。
    1N\frac{1}{N}

自动微分

训练神经网络，本质上是要回答一个问题：菜谱哪里该调整？传统做法靠厨师经验，机器学习靠梯度。在 JAX 里，求梯度不需要你手写反向传播，而是请来一位“营养分析师”——jax.grad。它拿到你的损失函数（菜谱卡），不改变原卡，而是生成一张新的“梯度菜谱卡”：输入同样的参数和数据，输出的是损失对参数的偏导数。

为什么 grad 能这么神奇？恰恰因为它要求被处理的函数是纯函数。营养分析师分析一份菜谱，只需要看食材和做法，不需要关心厨师的心情、厨房的温度、上次做菜留下的锅巴——这些外部因素（副作用）只会干扰分析。对于纯函数 L(θ, x, y)，jax.grad 可以精确地计算出“参数 θ 动一点点，损失 L 会动多少”：

有了梯度卡，配合 jax.jit 把函数编译成底层机器码、用 jax.vmap 把批量维度自动向量化，JAX 就把“纯函数”这个抽象变成了强大的性能引擎。营养分析师不只会看一份菜谱——它能把菜谱复印到所有分店、并行分析全部菜品。整个训练循环，就是“分析、调整、复印、再分析”的循环，而每一轮都是纯粹、干净、可复现的。

∇θL=grad(L)(θ,x,y)\displaystyle \nabla_\theta \mathcal{L} = \text{grad}(\mathcal{L})(\theta, x, y)
营养分析师拿到损失菜谱卡 $\mathcal{L}$，用 grad 这张魔法纸复印出一张梯度卡：喂给它同样的参数和数据，它告诉你每个调料该加还是该减、力度多大。
  1. $\mathcal{L}$（损失函数）衡量预测与真实值差距的纯函数，是“被分析”的对象。
    L\mathcal{L}
  2. θ（可微参数）grad 默认对第一个参数求导，所以要把需要训练的参数放在函数第一位。
    θ\theta
  3. grad（变换）JAX 的高阶变换，接收函数返回新函数，本身不修改原函数。
    grad(⋅)\text{grad}(\cdot)
  4. $\nabla_\theta \mathcal{L}$（梯度）损失对参数的偏导数向量，指明参数更新的方向和大小。
    ∇θL\nabla_\theta \mathcal{L}

AI 应用场景

大规模分布式训练

阶段：模型训练与加速

解决的问题：传统框架中权重散落在各设备的模型对象里，多卡同步需要小心翼翼处理锁和状态拷贝；纯函数范式将参数显式传出，每张卡各持一份参数副本做前向和反向，梯度聚合后统一生成新参数，天然无锁、无竞态。

具体案例：在 8 张 GPU 上使用 jax.pmap 并行训练视觉 Transformer。每张卡独立执行同一份“菜谱卡”，前向和反向互不干扰，最后对梯度做 all-reduce，再函数式地更新参数快照。整个流程没有共享可变状态，扩展性瓶颈只在于通信带宽。

训练实验复现与回溯

阶段：实验管理与调试

解决的问题：PyTorch 中参数被原地修改，想回到第 1200 轮的状态只能靠定期保存 checkpoint；纯函数范式下，每一轮更新都产生全新的参数对象，保存一个引用就等于拍下一张不会褪色的照片，复现和回溯成本极低。

具体案例：调参团队发现第 3000 轮的模型在验证集上表现异常，直接加载当时保存的参数快照，配合当时的输入数据重新执行前向推理，几分钟内定位到是某一层数值溢出，而非随机性干扰。

模型服务无状态切换

阶段：模型部署与在线推理

解决的问题：在线服务要切换模型版本时，传统做法需要重启进程或迁移内存中的对象状态；纯函数范式下，一次请求 = (params, input)，服务进程完全不记忆任何模型状态，加载一份新参数就能瞬时完成版本切换。

具体案例：推荐系统在服务端同时持有两套参数快照，新模型上线前先以 10% 流量做 A/B 对比。由于推理函数是纯函数，两套参数可以并存于同一进程，互不污染，对比结果干净可信。

Python 演示

python95 行
# -*- coding: utf-8 -*-# 纯函数范式演示: 用显式参数传递 + 函数式更新训练一个 MLPfrom copy import deepcopyfrom random import Random  def relu(v):    return max(0.0, v)  def init_params(layers, seed):    rng = Random(seed)    params = []    for i in range(len(layers) - 1):        w = [[rng.uniform(-0.5, 0.5) for _ in range(layers[i + 1])]             for _ in range(layers[i])]        b = [0.0] * layers[i + 1]        params.append((w, b))    return params  def forward(params, x):    """纯函数: 只读 params 和 x, 不修改任何外部状态"""    h = x    for (w, b) in params[:-1]:        h = [relu(sum(h[j] * w[j][k] for j in range(len(h))) + b[k])             for k in range(len(b))]    w, b = params[-1]    return [sum(h[j] * w[j][k] for j in range(len(h))) + b[k]            for k in range(len(b))]  def loss_fn(params, x, y):    out = forward(params, x)    return sum(0.5 * (out[k] - y[k]) ** 2 for k in range(len(y)))  def numerical_grad(params, x, y, eps=1e-5):    """数值梯度: 利用纯函数可复现性, 扰动参数后重新计算损失"""    base = loss_fn(params, x, y)    grads = []    for li, (w, b) in enumerate(params):        gw = []        for j in range(len(w)):            row = []            for k in range(len(w[j])):                p2 = deepcopy(params)                p2[li][0][j][k] += eps                row.append((loss_fn(p2, x, y) - base) / eps)            gw.append(row)        gb = []        for k in range(len(b)):            p2 = deepcopy(params)            p2[li][1][k] += eps            gb.append((loss_fn(p2, x, y) - base) / eps)        grads.append((gw, gb))    return grads  def update(params, grads, lr):    """函数式更新: 返回新参数, 旧参数对象保持不可变"""    new_params = []    for (w, b), (gw, gb) in zip(params, grads):        nw = [[w[j][k] - lr * gw[j][k] for k in range(len(w[j]))]              for j in range(len(w))]        nb = [b[k] - lr * gb[k] for k in range(len(b))]        new_params.append((nw, nb))    return new_params  def train():    # XOR 问题: 输入两个二进制位, 输出异或结果    data = [        ([0.0, 0.0], [0.0]),        ([0.0, 1.0], [1.0]),        ([1.0, 0.0], [1.0]),        ([1.0, 1.0], [0.0]),    ]    params = init_params([2, 8, 1], seed=42)    for epoch in range(60):        total = 0.0        for x, y in data:            loss_before = loss_fn(params, x, y)   # 旧参数只读            grads = numerical_grad(params, x, y)            params = update(params, grads, lr=0.5)  # 生成新参数            total += loss_fn(params, x, y)        if epoch % 10 == 0:            print("epoch %d loss %.6f" % (epoch, total / 4.0))    for x, y in data:        pred = forward(params, x)        print("input %s => %.3f (expect %.0f)" % (x, pred[0], y[0]))  if __name__ == "__main__":    train()
目的：用纯 Python 模拟 JAX 纯函数范式的 MLP 训练——参数显式传递、函数式更新、旧参数永不被修改，让你在不安装 JAX 的情况下体会函数式深度学习的核心风格。
  1. params = init_params([2, 8, 1], seed=42)用固定种子创建初始参数，保证每次运行起点一致，这是可复现的第一步。
  2. def forward(params, x):前向传播是纯函数，只读取 params 和 x，不修改任何外部状态。
  3. loss_before = loss_fn(params, x, y)更新前先记录旧参数下的损失，旧参数在这个过程中只读不改。
  4. p2 = deepcopy(params)计算梯度时复制一份参数再扰动，原参数对象完全不受影响。
  5. params = update(params, grads, lr=0.5)函数式更新：右边用旧参数算出一份全新参数，左边变量名重新绑定，旧参数仍可回溯。
  6. if epoch % 10 == 0:周期性打印平均损失，观察纯函数式训练循环是否在稳定收敛。

常见误区

误区我以为纯函数范式就是消灭一切状态，模型变成无记忆的“一次性工具”，那训练还怎么保存学到的知识？

正解纯函数不是没有状态，而是把状态从“隐藏的角落”搬到“明面的参数”。模型权重、优化器动量、随机数 key 全部作为显式参数传入传出。状态依然存在，只是不再偷偷摸摸，而是像食材一样摆在餐桌上，任何人一眼就能看到这道菜用了什么。

为何常见：因为“无副作用”常被误解为“无状态”，其实副作用指的是对外界隐式的影响，而显式传入传出的状态完全合法。

误区我在 PyTorch 里习惯先 `torch.manual_seed(0)` 再采样，到了 JAX 也照猫画虎，结果每跑一次结果都不一样，太诡异了。

正解JAX 的随机数生成器本身是有状态的，而这个状态必须显式传递。正确做法是把 key 作为函数参数传入，每次用 jax.random.split(key) 分裂出新 key，旧 key 保持原样。只有这样才能保证“相同输入相同输出”的纯函数性质。

为何常见：全局种子是面向对象时代的习惯，它依赖生成器内部状态，而这正是纯函数范式要消除的隐藏状态。

误区每次训练迭代都要生成一份全新参数，那不是每步都要复制一遍几十亿参数的模型？这不得慢死？

正解JAX 的数组是不可变对象，但底层通过 XLA 编译和缓冲区复用，params = update(params, grads) 并不会真的逐元素复制。函数式更新表达的是“逻辑上的新对象”，编译器会在不破坏不可变语义的前提下复用内存，实际运行时非常高效。

为何常见：我们习惯了 Python 对象的深拷贝思维，但 JAX 的不可变更新有编译器级别的优化，不能拿普通 Python 列表的复制成本来类比。

误区写论文用 JAX 炫一下可以，但真实生产系统还是 PyTorch 的 nn.Module 靠谱，纯函数太抽象了。

正解纯函数范式在工程上反而有独特优势：天然无锁并行、精确复现、缓存友好、模块可独立测试。多家大厂已经在 JAX 上跑工业级推荐系统和大模型训练。范式本身不分“学术”和“工业”，只分“是否适合你的团队和场景”。

为何常见：面向对象是主流入门教学的心智模型，人们对不熟悉的东西天然缺乏安全感，容易把“不习惯”误判成“不实用”。

误区我只要在函数里不写 `self.`，把所有变量都当作局部变量，就算“函数式”了吧？

正解真正的纯函数范式是一整套约定：参数必须显式传递、可变状态必须不可变更新、随机数必须用 key 管理、副作用必须彻底清除。只改表面写法，仍然可能偷偷依赖全局变量或外部状态。写 JAX 需要的是换一套心智模型，而不是换一套语法糖。

为何常见：语言语法相似会让人低估思维范式的差异，就像换了菜谱的排版格式，但厨师还是那个乱抓调料的厨师。

面试题

Q1. 下列对纯函数的描述，正确的是哪一项？

单选题 · easy

  • A.纯函数可以修改全局变量来记录调用次数
  • B.相同输入必然产生相同输出且无副作用← 正确答案
  • C.纯函数可以依赖随机数生成器的内部状态
  • D.纯函数输出与对象的内部属性和历史有关

答案：B。纯函数的两条铁律是“相同输入必得相同输出”和“无副作用”。选项 A 修改全局变量是副作用，C 依赖生成器内部状态会让输出不稳定，D 依赖对象内部属性属于隐藏状态，都与纯函数定义相悖。

Q2. 在 JAX 纯函数范式中，模型参数通常如何传递？

单选题 · mid

  • A.通过全局变量在模块间共享模型权重
  • B.将权重保存在类的实例属性中随时访问
  • C.以 Pytree 结构作为函数参数显式传递← 正确答案
  • D.使用进程间共享内存传递参数状态

答案：C。JAX 要求所有可变状态显式出现在函数签名中，参数打包成 Pytree 结构作为一个整体传入。全局变量、实例属性和共享内存都是把状态藏在函数外部，违背纯函数原则。

Q3. JAX 中随机数应如何管理才能保持纯函数性？

单选题 · mid

  • A.设置全局随机种子后直接调用随机采样
  • B.每次用 key 作为参数传入并分裂新 key← 正确答案
  • C.用系统时间为种子保证每次结果不同
  • D.把随机种子写死在函数内部即可复用

答案：B。全局种子依赖生成器的内部状态，系统时间引入外部不确定因素，种子写死在函数内则无法灵活控制。JAX 的做法是把 key 作为显式参数传入，采样后用 split 分裂出新 key，保证每次调用都是可复现的纯函数。

Q4. 纯函数范式带来的工程优势包括哪些？

多选题 · mid

  • A.相同输入输出稳定确保训练可复现← 正确答案
  • B.无隐藏状态天然适配并行与分布式← 正确答案
  • C.自动降低了所有矩阵运算的计算量
  • D.调试时可通过输入输出快照回溯问题← 正确答案

答案：ABD。可复现、易并行、可回溯都是纯函数无副作用的直接结果。但纯函数并不会自动降低矩阵运算量，计算复杂度取决于算法本身，与函数是否纯粹无关，因此 C 错误。

Q5. 与 PyTorch nn.Module 相比，JAX 纯函数范式的核心区别是什么？

单选题 · easy

  • A.参数显式作为函数输入而非对象属性← 正确答案
  • B.JAX 框架内部不支持神经网络的训练
  • C.PyTorch 无法计算损失函数的梯度值
  • D.两种范式在代码写法上没有任何区别

答案：A。PyTorch 把参数保存在 self 属性中由 model(x) 隐式读取，JAX 则把参数 θ 与输入 x 并列显式传入 f(θ, x)。B、C 均为错误陈述，D 与事实相反，两者在状态管理和心智模型上差异明显。

Q6. 关于 jax.grad 应用于纯函数，下列说法正确的有哪些？

多选题 · hard

  • A.grad 默认只对函数的第一个参数求导← 正确答案
  • B.要求被求导函数是消除副作用的纯函数← 正确答案
  • C.grad 会原地修改原函数的参数值
  • D.grad 返回新函数并支持和其他变换组合← 正确答案

答案：ABD。jax.grad 默认对第一个参数求导，所以可微参数要放首位；它要求函数是纯函数，否则无法精确分析；它返回一个新函数而不修改原函数，且可与 jit、vmap 组合。C 说反了，函数式范式中旧参数永远不会被原地修改。

Q7. 在函数式更新 params = update(params, grads) 中，旧参数对象发生了什么？

单选题 · extreme

  • A.旧参数对象被原地覆盖以节约内存空间
  • B.旧参数保持不变新参数是全新对象结构← 正确答案
  • C.旧参数立刻被回收且无法再被任何引用
  • D.只有开启缓存模式时旧参数才不改变

答案：B。函数式更新的核心是不可变性：右边的 params 作为输入只读，update 基于它构建并返回全新的参数结构，旧对象依然存在。选项 A 是原地修改的面向对象做法，C 忽略了旧参数仍可被其他引用持有可能，D 与是否开启缓存无关。

Q8. jax.vmap 变换的主要作用是什么？

单选题 · mid

  • A.将函数一次性编译为底层机器码显著加速
  • B.自动把批量维度向量化映射进目标函数← 正确答案
  • C.将纯函数转换为有状态的模块化类对象
  • D.自动保存训练断点并支持任务的恢复

答案：B。vmap 自动把批量维度映射进函数内部，让你无需手动写 for 循环就能逐样本调用纯函数，同时保持向量化性能。A 描述的是 jit，C 与纯函数方向相反，D 是 checkpoint 的功能，均与 vmap 无关。

Q9. 关于纯函数模块化设计，下列说法错误的是哪一项？

单选题 · hard

  • A.小纯函数可通过组合构建出大型模型
  • B.高阶函数可以把变换当作参数传递
  • C.纯函数模块易于独立测试和复用
  • D.模块化设计强制要求使用类与继承← 正确答案

答案：D。函数式模块化靠的是函数复合和高阶函数，并不需要类与继承——D 是错误说法。A、B、C 都是纯函数模块化的正确描述：小函数可组合、变换可作为参数、无副作用让测试与复用变得简单。

Q10. 下列哪些行为违反了纯函数原则？

多选题 · easy

  • A.在函数内打印日志并写入全局文件← 正确答案
  • B.在函数内修改传入的列表参数元素← 正确答案
  • C.函数返回输入经过计算后的全新结果
  • D.函数依赖系统当前时间决定输出值← 正确答案

答案：ABD。写文件是副作用，修改传入参数是原地篡改，依赖系统时间会让相同输入产生不同输出，三者都违反纯函数原则。C 返回新计算结果正是纯函数的典型行为，不违反任何规则。