想象你是一家黑科技餐厅的后厨总管。今晚有 100 位客人同时点了同一道招牌菜，菜单上还贴着一张“越做越好吃”的挑战书。你手下的三位大厨各怀绝技：品菜大厨只要尝一口成品，就能立刻说出每种调料该加还是该减、加多少——菜品的“改进方向”尽在掌握；备菜指挥不慌不忙，早在营业前就把所有流程预演优化成一条流水线，真正开火时绝不临时翻菜谱；批量出餐快手则一挥手点燃 100 个灶台，同一道菜同时出锅，绝不一份份慢慢等。这三位，正是 JAX 框架的核心三大 API：grad（自动微分）、jit（即时编译）、vmap（自动向量化）。JAX 把这一切装进一间“纯函数厨房”——每道菜都有标准菜谱、不依赖任何隐藏的老汤，于是三位大厨才能放心施展绝技。这篇文章，就带你走进这间厨房，看看三位大厨如何各司其职，又如何联手端出一盘惊艳的“可微分高速训练大餐”。

岗位关联

掌握 grad/jit/vmap 是踏入 JAX 生态的敲门砖，也是做深度学习研究、科学计算与元学习等高效可微分计算的职业基本功。

目录
  1. 可视化演示
  2. JAX框架初识
  3. grad自动微分
  4. jit即时编译
  5. vmap向量化
  6. 三者组合
  7. 注意事项
  8. AI 应用场景
  9. Python 演示
  10. 常见误区
  11. 面试题

可视化演示

\u53ef\u89c6\u5316\u5c55\u793aJAX\u6846\u67b6\u6838\u5fc3\u53d8\u6362\uff1agrad\u81ea\u52a8\u5fae\u5206\u3001jit\u5373\u65f6\u7f16\u8bd1\u52a0\u901f\u3001vmap\u81ea\u52a8\u5411\u91cf\u5316\u6279\u5904\u7406\uff0c\u4ee5\u53ca\u4e09\u8005\u7ec4\u5408\u5d4c\u5957\u6784\u5efa\u9ad8\u6548\u53ef\u5fae\u5206\u8ba1\u7b97\u7ba1\u7ebf\u3002

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

grad\u81ea\u52a8\u5fae\u5206\u53d8\u6362 \u51fd\u6570\u53d8\u6362\u51fd\u6570\u8303\u5f0f\uff0c\u63a5\u6536\u6807\u91cf\u51fd\u6570\u8fd4\u56de\u68af\u5ea6\u51fd\u6570 \u94fe\u5f0f\u6c42\u5bfc\u81ea\u52a8\u6784\u5efa\u8ba1\u7b97\u56fe f(x) \u2192 grad(f) \u2192 \u2207f \u70b9\u51fb\u5207\u6362\u67e5\u770bjit\u5373\u65f6\u7f16\u8bd1

JAX框架初识

每晚营业前，你这位后厨总管都要清点一遍菜谱。你定的规矩很简单：一道菜，原料进去，成品出来，过程中不许偷用冰箱里的“神秘老汤”。这份菜谱，在 JAX 里就是一个“纯函数”——输入决定输出，没有任何隐藏状态，也绝不偷偷修改锅铲的位置。正因为菜谱如此“透明”，后厨里的一切工作才能被精确预测和升级。

JAX 是谷歌开源的数值计算框架，它的气质就是这间“纯函数厨房”。它把 NumPy 风格的数组运算与一门核心哲学结合起来：函数式自动微分。在 JAX 里，几乎一切操作都是“输入一个函数，返回一个变换后的新函数”。数组运算 jnp.dot(x, w) 只是最基本的食材处理；真正让 JAX 特别的是三位大厨——grad（自动微分）、jit（即时编译）、vmap（自动向量化）。它们都是“函数变换器”：不修改你的菜谱本身，只在你调用时给菜谱附加超能力。

从数学上看，一个 JAX 函数就是一个从向量到向量的映射。理解这个基本模型，是看懂后面所有变换的前提。

所以，当你把一份“纯函数菜谱”交给后厨时，三位大厨都能放心接手——因为菜谱不藏秘密，谁都可以在上面做文章。接下来，我们逐一拜访这三位大厨。

y=f(x),x∈Rn,y∈Rm\displaystyle y = f(x), \qquad x \in \mathbb{R}^{n}, \quad y \in \mathbb{R}^{m}
这个公式说明：JAX 眼中的一切计算，都是“输入一个向量、输出一个向量”的纯映射，就像菜谱把原料变成成品。
  1. 输入向量 $x$模型或计算流程的输入，比如一张图片展平后的像素数组，维度为 $n$。
    x∈Rn\displaystyle x \in \mathbb{R}^{n}
  2. 输出向量 $y$函数处理后的结果，可能是预测分数或特征向量，维度为 $m$。
    y∈Rm\displaystyle y \in \mathbb{R}^{m}
  3. 映射关系 $f$一条从 $n$ 维空间通向 $m$ 维空间的通道，对应你写出的普通 Python 函数。
    f:Rn→Rm\displaystyle f: \mathbb{R}^{n} \to \mathbb{R}^{m}

grad自动微分

现在登场的是第一位大厨——品菜大厨。他练就了一项绝技：不管什么菜，只要端到他面前尝一口，他就能告诉你，这道菜里的每一种调料该往哪个方向调、调多少，菜才会更美味。在机器学习里，这道菜就是损失函数，而“调料的调整方向和幅度”，正是梯度。

grad 是 JAX 的自动微分 API。它把“求梯度”这件事变成了一个函数变换：你给它一个计算损失 $\mathcal{L}$ 的函数，它返回一个新的函数，这个新函数专门计算原函数对第一个参数的梯度。整个过程无状态、无副作用——品菜大厨不会替你炒菜，只给你一张“怎么调”的新菜谱。

自动微分的数学支柱是链式法则。品菜大厨之所以能一口尝出全盘问题，是因为他悄悄把每一道工序都“尝”了一遍：从出锅的菜往前追溯，一层层记录每个环节对最终味道的影响，最后合成完整的调整方案。

更妙的是，品菜大厨还能培养徒弟——grad(grad(f)) 可以求二阶导数，相当于培养出一位“品尝味道变化速度”的资深美食评论家。这在元学习、最优化和科学计算中都是利器。

∇θL(θ)=∑i∂L∂zi⋅∂zi∂θ\displaystyle \nabla_{\theta} \mathcal{L}(\theta) = \sum_{i} \frac{\partial \mathcal{L}}{\partial z_{i}} \cdot \frac{\partial z_{i}}{\partial \theta}
这个公式描述的是：参数 $\theta$ 的最终梯度，等于沿着计算链把所有中间环节的“影响系数”乘起来再求和。
  1. 中间变量 $z_i$损失函数内部每一层计算产生的中间结果，好比烹饪中的每个半成品。
    zi=gi(θ)\displaystyle z_{i} = g_{i}(\theta)
  2. 局部梯度 $\frac{\partial \mathcal{L}}{\partial z_i}$最终损失对第 $i$ 个中间结果的敏感度，即“这个半成品的好坏对最终味道影响多大”。
    ∂L∂zi\displaystyle \frac{\partial \mathcal{L}}{\partial z_{i}}
  3. 参数梯度 $\frac{\partial z_i}{\partial \theta}$第 $i$ 个中间结果对参数 $\theta$ 的敏感度，即“调料变化如何改变这个半成品”。
    ∂zi∂θ\displaystyle \frac{\partial z_{i}}{\partial \theta}
  4. 链式求和把所有路径的影响加起来，就得到损失对每个参数的梯度——这正是品菜大厨的“全盘调整方案”。
    ∇θL(θ)=∑i∂L∂zi⋅∂zi∂θ\displaystyle \nabla_{\theta} \mathcal{L}(\theta) = \sum_{i} \frac{\partial \mathcal{L}}{\partial z_{i}} \cdot \frac{\partial z_{i}}{\partial \theta}

jit即时编译

第二位大厨是备菜指挥。他的理念是：真正的效率不是开火后手忙脚乱，而是营业前把所有流程布局妥当。他会把你的菜谱通读一遍，在脑子里把每把刀、每个灶台的位置都安排到最优，再把流程浓缩成一条近乎肌肉记忆的流水线。等正式出餐时，他只需一声令下，所有步骤行云流水。

jit 就是这位备菜指挥。它用 XLA（Accelerated Linear Algebra，加速线性代数）把你的 JAX 函数提前编译成高效的机器代码。首次调用时，它默默完成“预演”（编译）；之后每一次调用都直接按流水线开火，通常能带来十倍以上速度提升。对于像 y = Wx + b 这样的密集矩阵运算，编译优化尤其显著。

编译是按“参数形状 + 类型”建立缓存的。如果传入同样形状的数组，第二次调用就直接命中缓存，连“预演”都省了。这就像一个后厨记住每位熟客的口味偏好——同一个熟客再来，无需重新询问。因此，jit 特别适合训练循环里“每步参数形状不变”的稳定场景。

不过，备菜指挥有个原则：菜谱绝不能临时改。如果你的函数里偷偷打印、读取全局变量，他会按“预演那一刻”的快照理解你的菜谱，结果往往与预期不符。这提醒我们：jit 只信任纯函数。

y=Wx+b,W∈Rm×n\displaystyle y = W x + b, \qquad W \in \mathbb{R}^{m \times n}
这是被编译优化的目标计算：一个典型的线性层，备菜指挥会把它打磨成最顺滑的机器指令序列。
  1. 权重矩阵 $W$$m \times n$ 的参数矩阵，决定输入每个分量如何组合。
    W∈Rm×n\displaystyle W \in \mathbb{R}^{m \times n}
  2. 输入 $x$ 与偏置 $b$$x$ 是 $n$ 维输入，$b$ 是 $m$ 维偏置项。
    x∈Rn,b∈Rm\displaystyle x \in \mathbb{R}^{n}, \quad b \in \mathbb{R}^{m}
  3. 矩阵乘加核心运算是矩阵乘法再叠加偏置，备菜指挥会把这个运算编译成高效的机器码循环。
    y=Wx+b\displaystyle y = W x + b

vmap向量化

第三位大厨是批量出餐快手。他负责处理最折磨人的场景：100 位客人同时点了同一道菜。普通厨师会一份一份炒，他却能一声吆喝点燃 100 个灶台，几十秒全部出锅。味道一模一样，时间却省了几十倍。这正是深度学习中“批量训练”的痛点：处理一个样本和处理一万个样本，不应该写成一遍又一遍的循环。

vmap 就是这个能力：自动向量化。你写好处理“单个样本”的函数，vmap 自动在它前面加一个“批量维度”，让函数同时处理一批样本。比如 100 张图片，不用手写 for 循环，一个 vmap 全搞定。它把“逐样本逻辑”和“批量执行”彻底解耦，代码既简洁又不容易出错。

更优雅的是，vmap 允许通过 in_axes 指定哪些参数被批量映射、out_axes 指定输出批量维度放在哪里。它还可能触发 XLA 的批处理优化，比手写 Python 循环快一大截——相当于 100 个灶台之间还有协同管道，共享预热和调料，效率远超各自为战。

批量出餐快手最拿手的是协同作战：当它和品菜大厨搭档时，能一口气尝出 100 份菜各自的调整方向——这正是批量梯度下降的雏形。

vmap(f)([x1,x2,…,xB])=[f(x1),f(x2),…,f(xB)]\displaystyle \text{vmap}(f)([x_{1}, x_{2}, \ldots, x_{B}]) = [f(x_{1}), f(x_{2}), \ldots, f(x_{B})]
这个公式表明：vmap 把“对单个样本的函数”升级为“对整个批量的同时处理”，就像把单人份菜谱升级成百人宴。
  1. 单样本函数 $f$你只需写出处理一个样本的逻辑，比如单张图片的预测。
    f:Rn→Rm\displaystyle f: \mathbb{R}^{n} \to \mathbb{R}^{m}
  2. 批量输入列表$B$ 个样本拼成一个批量数组，第一维就是批量维度。
    [x1,x2,…,xB],xi∈Rn\displaystyle [x_{1}, x_{2}, \ldots, x_{B}], \quad x_{i} \in \mathbb{R}^{n}
  3. 逐样本映射vmap 把 $f$ 自动应用到每个样本上，输出拼成批量结果。
    [f(x1),f(x2),…,f(xB)]\displaystyle [f(x_{1}), f(x_{2}), \ldots, f(x_{B})]

三者组合

三位大厨各怀绝技，但真正的招牌菜，是让他们联手上灶。顾客说新菜品还不够完美，于是品菜大厨负责尝味找方向，批量出餐快手负责一次试吃 100 份样品，备菜指挥则把整个流程预演到最快。三位一体，一道菜在几分钟内就完成了从“能入口”到“惊艳全场”的进化。

在 JAX 中，组合就是用函数变换“套娃”。最经典的高效训练管线是：先用 vmap 把单样本损失变成批量损失，再用 grad 求批量损失对参数的梯度，最后用 jit 把整个步骤编译加速。三者各有分工，又不互相干扰，正是函数式设计的魅力。

组合顺序值得玩味：jit(grad(vmap(f))) 是常见写法——先向量化，再求导，最后编译。每一层变换都是纯函数操作，所以顺序可以灵活调整，不改变数学结果，只影响执行性能。这就像让三位大厨任意排班，不论谁先谁后，最终端出的菜都一样好吃，只是出餐速度不同。

于是，三位大厨的合奏形成了一条完整的训练管线：尝味、批量试吃、备菜预演，三位一体，缺一不可。

θ←θ−η⋅1B∑i=1B∇θL(fθ(xi),yi)\displaystyle \theta \leftarrow \theta - \eta \cdot \frac{1}{B} \sum_{i=1}^{B} \nabla_{\theta} \mathcal{L}(f_{\theta}(x_{i}), y_{i})
一次完整的批量参数更新：尝遍 100 份样品，平均出总体调整方向，再沿该方向迈出一小步。
  1. 批量样本损失每个样本都算出一个损失值，$\mathcal{L}(f_{\theta}(x_i), y_i)$ 是第 $i$ 个样本的误差。
    L(fθ(xi),yi)\displaystyle \mathcal{L}(f_{\theta}(x_{i}), y_{i})
  2. 逐样本梯度对每个样本的损失分别求参数梯度，相当于 100 份菜各自的调整建议。
    ∇θL(fθ(xi),yi)\displaystyle \nabla_{\theta} \mathcal{L}(f_{\theta}(x_{i}), y_{i})
  3. 平均梯度把所有样本的梯度求平均，得到这一批的总体调整方向。
    1B∑i=1B∇θL(fθ(xi),yi)\displaystyle \frac{1}{B} \sum_{i=1}^{B} \nabla_{\theta} \mathcal{L}(f_{\theta}(x_{i}), y_{i})
  4. 参数更新沿着负梯度方向迈出学习率为 $\eta$ 的一步。
    θ←θ−η⋅1B∑i=1B∇θL(fθ(xi),yi)\displaystyle \theta \leftarrow \theta - \eta \cdot \frac{1}{B} \sum_{i=1}^{B} \nabla_{\theta} \mathcal{L}(f_{\theta}(x_{i}), y_{i})

注意事项

厨房再好，规矩也要守住。三位大厨共用一个前提：你的菜谱必须是“纯函数”。所谓纯函数，就是同样的原料一定做出同样的菜，绝不依赖或修改外界状态。一旦菜谱偷偷用了“冰箱里的老汤”，大厨们就会看走眼——grad 会算错方向，jit 会缓存过期结果，vmap 会得到错误批量输出。

最常见的坑有三个：一是 jit 和 grad 内部不能用依赖动态形状的 Python for，也不能有 print 等副作用；二是闭包中依赖外部计数器时，追踪期只会记录初值；三是调试时被编译打包的函数难以逐行查看。比如想统计每次前向传播的调用次数，直接在函数外写 count += 1 是无效的，因为追踪阶段只执行一次。

JAX 也给了对策：动态控制流请用 jax.lax.cond / jax.lax.while_loop；统计用的计数变量放进一个列表容器中，用 list.append 记录；调试可以用 jax.debug.print，或先不套 jit 直接跑小样例。这些技巧像是给大厨们配了一副“透视眼镜”，让纯函数的边界清晰可见。

守住纯函数这条红线，三位大厨就能放心施展。下一次，当你把一份“带隐形状态”的菜谱递给 JAX 时，不妨先想想：这道菜在我的后厨里，真的经得起大厨们的“直球检验”吗？

f=g⟹jit(grad(f))=jit(grad(g))\displaystyle f = g \quad \Longrightarrow \quad \text{jit}(\text{grad}(f)) = \text{jit}(\text{grad}(g))
纯函数的安全性表现为“变换保持等价”：两个行为一致的菜谱，经过大厨改造后依然一致。
  1. 函数等价 $f = g$对相同的输入，两个函数给出完全相同的输出。
    f=g\displaystyle f = g
  2. 变换保持grad、jit、vmap 等函数变换不会破坏这种等价性。
    grad(f)=grad(g)\displaystyle \text{grad}(f) = \text{grad}(g)
  3. 组合不变性多重变换叠加后，等价性依然成立，这是纯函数世界最安心的“合同”。
    jit(grad(f))=jit(grad(g))\displaystyle \text{jit}(\text{grad}(f)) = \text{jit}(\text{grad}(g))

AI 应用场景

深度学习模型训练

阶段：模型训练与优化

解决的问题：传统训练循环要手写反向传播、手动管理 GPU 批处理，代码冗长且易错；JAX 用 grad + jit + vmap 把梯度计算、编译加速、批量处理全部声明式化，训练循环简洁且天然可微分。

具体案例：训练一个线性回归或小型 MLP 分类器时，用 jit(vmap(grad(loss))) 计算批量梯度并更新参数，代码量比传统方式少一半，速度却更快。本文 Python 演示就是这样一个完整可运行的例子。

科学计算与物理仿真

阶段：数据建模与参数反演

解决的问题：物理仿真常需对模拟器求导（如逆问题、参数标定），传统有限差分耗时且精度差；JAX 的 grad 可直接对仿真函数求精确梯度，配合 vmap 并行仿真多组初始条件。

具体案例：用 JAX 写一个可微分的流体或刚体模拟器，用 vmap 同时仿真 100 组初始条件，用 jit 加速迭代，最后用 grad 反演未知的初始参数，让“尝味大厨”直接作用于物理世界。

元学习与超参数搜索

阶段：模型选择与自动调优

解决的问题：元学习需要二阶梯度（对学习率或初始参数求导），手推公式痛苦不堪；JAX 的 grad(grad(f)) 可以直接嵌套求高阶导，把整个双层优化变成纯函数变换。

具体案例：在 MAML 元学习算法中，内环用 grad 更新每个任务上的参数，外环再用 grad 计算对初始参数的二阶梯度，最后套上 jit 加速整个双层循环，让“品菜大厨的徒弟”也上阵帮忙。

Python 演示

python41 行
import jaximport jax.numpy as jnpfrom jax import grad, jit, vmap  def predict(params, x):    w, b = params    return w * x + b  def loss_fn(params, x, y):    pred = predict(params, x)    return 0.5 * jnp.mean((pred - y) ** 2)  @jitdef train_step(params, x, y, lr):    grads = grad(loss_fn)(params, x, y)    w, b = params    gw, gb = grads    return (w - lr * gw, b - lr * gb)  batch_grad = vmap(grad(loss_fn), in_axes=(None, 0, 0))  def main():    xs = jnp.linspace(0.0, 1.0, 100)    ys = 2.0 * xs + 1.0    params = (jnp.array(0.0), jnp.array(0.0))    lr = 0.1    for epoch in range(300):        params = train_step(params, xs, ys, lr)        if epoch % 100 == 0:            gs = batch_grad(params, xs, ys)            loss = loss_fn(params, xs, ys)            print("epoch", epoch, "loss", float(loss), "gw_mean", float(gs[0].mean()))  if __name__ == "__main__":    main()
目的：用 JAX 三大 API 训练一个线性回归模型，演示 grad 求梯度、jit 编译加速、vmap 批量求梯度的完整组合。
  1. from jax import grad, jit, vmap一次导入三位大厨，本文主角全部登场。
  2. def predict(params, x)写出处理单个样本的预测函数，保持纯函数风格。
  3. @jit给训练步套上编译加速，让每次参数更新都快人一步。
  4. grads = grad(loss_fn)(params, x, y)对损失函数求梯度，返回的参数梯度元组。
  5. batch_grad = vmap(grad(loss_fn), in_axes=(None, 0, 0))批量求梯度：参数不批量，x 和 y 按轴 0 批量。
  6. params = train_step(params, xs, ys, lr)整个训练管线一气呵成，三位大厨协同出餐。

常见误区

误区初学 JAX 时以为 `grad(f)` 会直接计算出梯度值并保存在模型对象里，像 PyTorch 那样调用 `loss.backward()` 后从 `param.grad` 取数。

正解grad 是函数变换：它接收函数 $f$，返回一个新的梯度函数。你调用这个新函数，才得到梯度数组。整个过程无状态、无副作用。好比品菜大厨不是替你把菜调好，而是递给你一张“如何调”的新菜谱。

为何常见：习惯命令式框架的人很难立刻适应函数式风格，以为梯度是“某处维护的状态”而非“一个函数的返回值”。

误区以为给任何函数加上 `@jit` 都会变快，包括里面有 print、有动态 for、有全局变量的函数。

正解jit 只信任纯函数。它用追踪（tracing）方式把 Python 代码编译成 XLA 图，副作用只发生在追踪期，动态形状会直接报错。备菜指挥只会预演标准菜谱，遇到“随机应变”的菜谱就会罢工或给出过期结果。

为何常见：其他框架的“加速装饰器”（如 numba）处理范围不同，加上 JAX 宣传语的诱惑，容易让初学者高估 jit 的适用范围。

误区认为 `vmap` 本质就是 Python 的 for 循环，只是写法简洁一点，不会带来本质加速。

正解vmap 通过批处理规则在底层融合运算，配合 XLA 可以避免 Python 循环的解释开销，并有机会触发更优的并行计划；当与 jit 组合时，批量维度被纳入编译优化。好比批量出餐快手不是把 100 份菜挨个做，而是同时点燃 100 个灶台。

为何常见：从表面 API 看 vmap 确实像内置的 map 函数，这种“语法糖”错觉很自然，容易低估它的底层优化能力。

误区觉得 JAX 的 `grad` 只是玩具，复杂模型或二阶导数必须自己手推公式或用别的库。

正解grad 支持任意深度的嵌套组合，grad(grad(f)) 直接给出二阶导，在 MAML 元学习、牛顿法优化、科学计算中是常用操作。品菜大厨不仅能尝出味道，还能尝出味道变化的趋势。

为何常见：大多数教程只演示 loss 的一阶梯度，很少展示高阶导的真实应用，导致初学者视野受限。

面试题

Q1. grad(f) 对函数 f(x) 求梯度后，返回的结果是？

单选题 · easy

  • A.一个具体的标量梯度数值
  • B.一个可调用的梯度函数← 正确答案
  • C.一份更新后的参数列表
  • D.一个自动微分的计算图

答案：B。grad 是函数变换，输入一个函数，返回一个新函数。调用新函数才得到梯度值，而不是直接给出数值。品菜大厨给你的是“如何调”的菜谱，不是已经调好的菜。

Q2. jit 的核心作用是什么？

单选题 · easy

  • A.自动计算函数梯度
  • B.编译优化加速执行← 正确答案
  • C.批量映射多个样本
  • D.自动分配计算设备

答案：B。jit 通过 XLA 把函数提前编译成高效机器码，首次调用后命中缓存，执行速度大幅提升。它不负责求导、不负责批处理，只负责“备菜预演”。

Q3. vmap 的核心作用是什么？

单选题 · easy

  • A.把逐样本循环变成批量运算← 正确答案
  • B.自动求解二阶偏导数
  • C.把函数提前编译成机器码
  • D.自动释放不用的显存

答案：A。vmap 自动向量化，让处理单个样本的函数能同时作用于一整批样本，相当于批量出餐快手同时点燃 100 个灶台。

Q4. jit 函数的参数中包含值会变化的 Python 整数，应如何标记？

单选题 · mid

  • A.通过 static_argnums 指定位置← 正确答案
  • B.通过 dynamic_argnums 指定位置
  • C.通过 static 参数关闭缓存
  • D.该情况无法通过 jit 支持

答案：A。static_argnums 告诉 jit 哪些参数是“静态”的（值会影响控制流但形状不变），改动该值会触发重新编译。备菜指挥需要知道有些参数会影响流程布局，必须显式告知。

Q5. 关于 grad 的输出要求，下列说法正确的是？

单选题 · mid

  • A.输出必须是标量（如损失函数值）← 正确答案
  • B.输出必须是二维矩阵形式
  • C.输出可以是任意 Python 对象
  • D.输出必须是整数类型结果

答案：A。grad 要求被求导函数的输出是标量，因为梯度定义为标量损失对参数的偏导。品菜大厨只能对“一道菜的整体味道”评价，不能对一桌子菜分别评价——那是 jacobian 的职责。

Q6. 下列哪种输入会让 vmap 抛出异常？

单选题 · mid

  • A.形状为 (B, n) 的批量数组
  • B.长度互不相同的 Python 列表← 正确答案
  • C.形状相同的元组型参数组
  • D.带有 None 标记的批量维度

答案：B。vmap 要求批量维度上的形状一致，长度不一的 Python 列表无法映射成规则的批量数组。批量出餐要求每份菜的量一致，否则无法同时开火。

Q7. 在 jit 化的函数中调用 print，会发生什么？

单选题 · mid

  • A.每次调用都会打印一次
  • B.只在编译追踪时打印一次← 正确答案
  • C.永远不会被打印出来
  • D.会直接导致编译报错

答案：B。jit 用追踪方式执行 Python 代码，print 是副作用，只发生在追踪（编译）阶段，运行时机器码里没有它。备菜指挥只在预演时听你口令，正式出餐时不会再听一遍。

Q8. 计算函数的二阶导数，正确的写法是？

单选题 · hard

  • A.调用 grad(f, order=2)
  • B.嵌套调用 grad(grad(f))← 正确答案
  • C.调用 jit(f).grad(2)
  • D.使用 vmap(f, hessian=True)

答案：B。JAX 的 grad 可以嵌套使用，grad(grad(f)) 返回二阶导函数。品菜大厨的徒弟也能培养，一层层“尝出味道变化的趋势”。其他写法都不是 JAX 的 API。

Q9. jit 的正确使用前提包括哪些？

多选题 · mid

  • A.函数内不使用外部可变全局变量← 正确答案
  • B.每次调用时参数形状保持一致← 正确答案
  • C.函数内允许任意 Python 副作用
  • D.函数中不能出现 jnp 数组运算

答案：AB。jit 要求纯函数（无外部可变状态）且形状稳定，这样编译缓存才能复用。它恰恰是 jnp 数组运算的加速对象，而不是禁止对象。备菜指挥只接受标准菜谱，且每次出餐规格一致。

Q10. 关于批量梯度更新，下列说法正确的有？

多选题 · hard

  • A.vmap(grad(loss)) 可得到逐样本梯度← 正确答案
  • B.grad(vmap(loss)) 会因为输出非标量而报错← 正确答案
  • C.实践中常用逐样本梯度的平均值更新参数← 正确答案
  • D.vmap 与 grad 无法在同一个函数上组合

答案：ABC。vmap(grad(loss)) 让每个样本各自求梯度，输出批量梯度；grad(vmap(loss)) 中 vmap 的输出是向量，grad 要求标量，因此会报错；实践中对批量梯度求平均后更新参数。批量出餐和尝味可以先后搭配，但顺序不同，效果和可行性也不同。