想象你是一家餐厅的后厨总管。宴会厅的订单是提前三天锁定的——10 桌、每桌 8 人、菜单固定。你可以提前批量采购、切配、排好出菜顺序，开席后整个厨房行云流水。这叫静态形状：所有信息提前到位，主厨（编译器）能从容优化。而散客区完全相反——客人随到随坐，一桌 3 人、一桌 9 人、中途还加座。你只能“兵来将挡”，每次都要重新调整火力、临时分配人手，虽然灵活，但出菜慢、还容易乱。这叫动态形状。 JAX 的形状管理，本质就是这家后厨的调度艺术。数组的形状（有多少行、多少列、多少通道）分为静态与动态两类：静态形状在编译时确定，让底层编译器 XLA 大展拳脚；动态形状在运行时变化，增加了灵活性，却要付出性能代价。这篇教程将带你认识形状推导规则、混合形状处理（padding + mask）以及 static_argnums 的用法，让你写出的 JAX 代码既灵活又高效——既能接得住“散客”，又不浪费“宴席”的优化红利。

岗位关联

对 AI 从业者而言，掌握 JAX 形状管理是从“代码能跑”迈向“模型跑得快”的关键一跃，直接决定你在训练大模型、部署推理服务时能否榨干硬件性能。

目录
  1. 可视化演示
  2. 形状基础概念
  3. 静态形状优势
  4. 动态形状代价
  5. 形状推导规则
  6. 混合形状处理
  7. JIT与静态参数
  8. 常见形状错误
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

通过形状对比、XLA编译流程和Padding策略三个视角，展示JAX中静态形状的编译优化优势与动态形状的处理方法。

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

形状基础概念

餐厅后厨的账本上，每一张桌子都记着“几位客人”。宴会厅的 10 人桌是提前锁定的，账本上永远写着 10，这叫静态形状；而散客区的桌子，账本上的数字随着客人入座不断变化——5 人、8 人、3 人……这叫动态形状。JAX 的数组就像这些桌子，每个数组都有一个“形状”，描述它有多少行、多少列、多少通道的数字组合。

在 JAX 里，形状是一个用元组表示的属性，例如一个 4×4 的矩阵形状是 (4, 4)。关键区别在于“何时知道”：静态形状在编译（jax.jit 追踪）时就已经完全确定，编译器可以把形状当作“已知常数”刻进机器码；动态形状则要等真正运行到那一刻才知道，编译器只能“两手一摊”，先按最稳妥的方案准备。这个时机差异，就是 JAX 性能分水岭的起点。

拿 shape 算子来说，它做的事情极其简单：把一个张量的每一维大小完整提取出来，形成元组。用数学语言写就是：

回到后厨，“形状”就是账本上的“桌位数量”。写 JAX 代码时，每构造一个数组都相当于在本子上记下一笔：静态的账本让主厨安心备菜，动态的账本让主厨时刻待命。这就是形状管理的第一课。

shape(x)=(d1,d2,…,dn),di∈Z≥0\displaystyle \text{shape}(\mathbf{x}) = (d_1, d_2, \ldots, d_n), \quad d_i \in \mathbb{Z}_{\geq 0}
  1. 形状算子shape(x) 返回张量 x 的每一维大小组成的元组，实际代码里就是 x.shape。
    shape(x)\text{shape}(\mathbf{x})
  2. 维度分量d_i 表示第 i 维的大小，比如 (128, 256) 中 128 是行数、256 是列数。
    d1,d2,…,dnd_1, d_2, \ldots, d_n
  3. 非负整数域Z_{\geq 0}^n 说明每个维度都是非负整数，维度可以是 0（空张量），但绝不可能是负数。
    Z≥0n\mathbb{Z}_{\geq 0}^n

静态形状优势

宴会订单提前三天确定 10 桌，主厨能提前三天开始备菜：批量采购食材、统筹火候、安排人手，把整个出菜流程编排得像一场交响乐。静态形状给 XLA 编译器的正是这种“提前量”——所有形状信息在编译期就摊在桌面上，编译器可以大刀阔斧地优化。

JAX 中静态形状意味着 XLA 在编译期掌握全部形状信息，不需要在运行时“看一眼”数据才知道形状。所有布局、对齐、kernel 选择都在编译期决定。具体来说，XLA 会遍历所有可用的算子实现（kernel），挑出在这个静态形状上执行最快的那一个；还能做内存布局优化（比如 NCHW 与 NHWC 的抉择）、算子融合（把多个 elementwise 算子合并，减少内核启动次数）以及内存的预先分配。

写成公式，静态形状下的最优执行时间可以表达为：

提前知道 10 桌，主厨才敢选最合适的铁锅、最省火的灶台。静态形状给 XLA 的正是这种“确定性红利”——一次编译，多次复用最优方案。

Trun=min⁡k∈Kcost(k,Sstatic)\displaystyle T_{\text{run}} = \min_{k \in \mathcal{K}} \text{cost}(k, S_{\text{static}})
  1. 候选 kernel 集合K 是所有可用算子实现的集合，就像厨房里不同尺寸的铁锅、砂锅、蒸锅。
    K\mathcal{K}
  2. 代价函数cost(k, S) 估算 kernel k 在这个形状上执行需要的时间。
    cost(k,Sstatic)\text{cost}(k, S_{\text{static}})
  3. 最优选择min 表示从所有候选中挑出代价最小的那个。
    min⁡k∈K\min_{k \in \mathcal{K}}
  4. 静态形状S_static 是编译期已知的完整形状，让代价搜索在编译时就能完成。
    SstaticS_{\text{static}}

动态形状代价

散客高峰期，每一桌人数都不一样，主厨没法批量备菜，只能反复调整火力：这桌刚起锅，那桌又加人。动态形状的灵活性背后，是实实在在的效率代价。JAX 为了编译，必须要么等运行时拿到形状再临时编译，要么退而求其次用保守的通用 kernel。

动态形状的主要代价有三层。第一层是重编译风险：每次形状变化，之前编译好的机器码就作废了，需要从头编译一次，GPU 在这段时间只能闲置。第二层是次优 kernel：编译器在编译期不知道形状，只能选择“什么形状都能跑”的通用实现，放弃针对特定形状的极致优化。第三层是运行时开销：每次调用都要检查形状、走动态分发（dispatch）逻辑，这些在静态形状下都可以省掉。

把这些代价写进一个公式，动态形状下的总时间可以分解为：

散客越多，主厨花在“重新安排”上的时间就越多，真正炒菜的时间反而变少了。这就是动态形状的“灵活之痛”——你获得了变长输入的自由，却要支付编译器的重复劳动。

Tdynamic=Tstatic+∑i=1NTrecomp(Si)⋅1[Si≠Si−1]\displaystyle T_{\text{dynamic}} = T_{\text{static}} + \sum_{i=1}^{N} T_{\text{recomp}}(S_i) \cdot \mathbb{1}[S_i \neq S_{i-1}]
  1. 理想成本T_static 是形状固定不变时最优执行时间，对应宴席场景的“完美出菜时间”。
    TstaticT_{\text{static}}
  2. 指示函数1[S_i ≠ S_{i-1}] 是一个开关，当第 i 次的形状和上一次不一样时取 1，否则取 0。
    1[Si≠Si−1]\mathbb{1}[S_i \neq S_{i-1}]
  3. 重编译代价T_recomp(S_i) 是为新形状 S_i 重新编译所花的额外时间。
    Trecomp(Si)T_{\text{recomp}}(S_i)
  4. 加总把每次形状变化引发的重编译时间全部累加起来。
    ∑i=1N\sum_{i=1}^{N}

形状推导规则

厨房里合并食材，一人份的料和四人份的料要想做成一大盘，必须按“对齐规则”把分量统一——这正是厨房里的广播规则。JAX 的广播机制，就是形状推导的核心规则：两个形状不同的数组做运算时，维度会自动扩展对齐。

广播规则只有三条：维度从右往左对齐比较；相同的维度直接通过；某个维度为 1 时扩展到另一个维度的大小。举例来说，形状 (3, 1) 的数组和 (1, 4) 的数组相加，从右往左对齐：第一维 1 与 4 取 4，第二维 3 与 1 取 3，最终得到 (3, 4) 的输出——就像 3 行“一人份”和 4 列“四人份”合并成 3×4 的“大拼盘”。

任意一个算子在编译期都会经过形状推导：输入形状映射到输出形状，这是纯函数运算，不依赖具体数值。规则的数学表达是：

主厨只要记住三句对齐口诀，就能在备菜前判断两批食材能不能合并成一道菜。广播规则就是 JAX 后厨的“对齐口诀”——它决定了数组之间能否协作、如何协作。

out_shapei={aiif ai=biaiif bi=1biif ai=1errorotherwise\displaystyle \text{out\_shape}_i = \begin{cases} a_i & \text{if } a_i = b_i \\ a_i & \text{if } b_i = 1 \\ b_i & \text{if } a_i = 1 \\ \text{error} & \text{otherwise} \end{cases}
  1. 从右往左对齐比较两个形状的第 i 维，i 从最右边（最后一维）开始。
    ai,  bia_i, \; b_i
  2. 相等则通过两维相等时，输出取这个相同的值，无需扩展。
    ai=bi⇒out_shapei=aia_i = b_i \Rightarrow \text{out\_shape}_i = a_i
  3. 一者为 1 则扩展某个维度为 1 时，按另一个维度的值扩展，相当于“一人份复制成多人份”。
    bi=1⇒out_shapei=aib_i = 1 \Rightarrow \text{out\_shape}_i = a_i
  4. 都不匹配则报错既不相等也非 1，广播失败，抛出 shape mismatch 错误。
    otherwise⇒error\text{otherwise} \Rightarrow \text{error}

混合形状处理

散客桌数不定，聪明的主厨不纠结于每桌实际人数，而是直接按最大桌数摆好长桌，暂时没人的座位放一个“已预订”牌——真实客人落座的位置才上菜。这就是 padding + mask 的智慧：先让形状统一，再让计算忽略无效位。

padding 是处理变长输入最实用的手段：把所有序列填充到 batch 内的最大长度，让一切形状变得静态、统一。但填充出来的“假数据”不能参与真正的计算，于是引入 mask——一个与填充后形状相同的布尔矩阵，标记哪些位置是真实数据。在 NLP 中，这几乎是标配：一句话 3 个词、另一句 30 个词，统一 pad 到 30，再用 mask 告诉模型“第 3 个词之后都是占位符”。

具体来说，给第 i 个序列的真实长度 L_i，填充到 L_max 后形状变为 (B, L_max, E)。mask 矩阵 M 中 M[i][j] = 1 当且仅当 j < L_i。计算 loss 时用 mask 把填充位置的权重置零，再除以真实长度，填充位就不会影响梯度。数学表达如下：

“已预订”牌让主厨既能按统一规格备菜，又不会给空位真的上菜。padding + mask 让 JAX 既享受静态形状的编译优化，又不丢失变长数据的灵活性——这是处理动态形状最经典的“两条腿走路”方案。

X^=pad(X,(0,Lmax⁡−Li)),Mij=[j<Li]\displaystyle \hat{X} = \text{pad}(X, (0, L_{\max} - L_i)), \quad M_{ij} = [j < L_i]
  1. pad 操作pad(X, (0, L_max - L_i)) 把长度 L_i 的序列在末尾补零到 L_max。
    pad(X,(0,Lmax⁡−Li))\text{pad}(X, (0, L_{\max} - L_i))
  2. 填充后的形状X_hat 形状为 (B, L_max, E)，所有序列整齐划一。
    X^∈RB×Lmax⁡×E\hat{X} \in \mathbb{R}^{B \times L_{\max} \times E}
  3. 掩码矩阵M[i][j] = 1 表示第 i 个序列的第 j 个位置是真实数据。
    Mij=[j<Li]M_{ij} = [j < L_i]
  4. 忽略填充位计算时将 M 为 0 的位置权重清零，相当于“已预订”牌挡住空位。
    loss=∑Mij⋅ℓij∑Mij\text{loss} = \frac{\sum M_{ij} \cdot \ell_{ij}}{\sum M_{ij}}

JIT与静态参数

菜单上有些“隐藏菜”只按固定规格出品，比如 10 人份的佛跳墙，主厨必须提前三天准备。static_argnums 就是你在下单时告诉主厨：“这道菜就按 10 人份做，别改了。” 这样主厨就能提前锁定制法、备料、排期。

jax.jit 默认把所有参数当作动态值（tracer），形状在编译期是“未知数”。但有些参数本质上是“配置”而非数据——比如模型隐层维度、循环次数、整数标志位。把它们标记为 static_argnums，参数就变成编译期常量，可以参与 Python 切片、控制流，让编译器看到完整形状。

调用 jax.jit(f, static_argnums=(1,)) 后，第 1 个参数（索引从 0 开始，所以是第二个参数）必须是可哈希的 Python 值（如 int、tuple）。每次该参数值变化都会触发重新编译，生成独立版本。所以它适合“稳定且影响形状”的参数，绝不适合“每个 batch 都变”的参数——否则你会陷入无休止的重编译。数学表达：

固定规格的佛跳墙出菜快，但客人一改规格就得重新准备。static_argnums 是把双刃剑——用对是预编译优化，用错是重编译灾难。判断标准很简单：这个参数真的稳定吗？它影响形状或控制流吗？

fjit=jax.jit(f,static_argnums=(1,))\displaystyle f_{\text{jit}} = \text{jax.jit}(f, \text{static\_argnums}=(1,))
  1. jit 包装f_jit 是函数 f 经过 JIT 编译后的版本，调用它自动命中编译缓存。
    fjit=jax.jit(f)f_{\text{jit}} = \text{jax.jit}(f)
  2. 静态参数指定static_argnums=(1,) 告诉编译器第二个参数是编译期常量。
    static_argnums=(1,)\text{static\_argnums}=(1,)
  3. 编译差异静态参数值不同，生成的编译版本就不同，f_jit(x, 32) 与 f_jit(x, 64) 是两份独立机器码。
    fjit(x,32)≠fjit(x,64)f_{\text{jit}}(x, 32) \neq f_{\text{jit}}(x, 64)

常见形状错误

后厨最怕“上错菜”——客人点的是 4 人份，厨房却端来 3 人份。要找出错在哪，得先核对每张桌的“人数清单”（形状元组），再从上菜流程（计算图）里逐环节排查。JAX 中的形状错误同样有章可循。

最常见的错误是 shape mismatch：两个张量做运算时，维度从右往左对齐后发现既不相同、也没有一个为 1，于是广播失败，抛出 ValueError。另一个高发错误是 ConcretizationTypeError（具体化错误）：在 jit 函数里把动态 tracer 当作具体 Python 值用，比如 if x.shape[0] > 5 或 for i in range(x.shape[0])——tracer 的值在编译期不可知，Python 的分支与循环需要一个“具体的数”。

排查 shape mismatch 的利器就是广播规则本身：任一维度不满足相等或为 1 的条件即报错。判定条件写清楚就是：

核对人数清单、逐桌排查，才能快速定位上错菜的环节。排查具体化错误也有技巧：用 jax.debug.print 在编译期打印 tracer 信息，或把关键整数参数放进 static_argnums。记住：JAX 的控制流条件必须是具体值，动态条件请用 jax.lax.cond。形状错误的排查不是瞎猜数字，而是对照确定性的规则逐维核验。

error  ⟺  ∃i:ai≠bi∧ai≠1∧bi≠1\displaystyle \text{error} \iff \exists i: a_i \neq b_i \land a_i \neq 1 \land b_i \neq 1
  1. 存在量化存在 i 表示只要有一维出问题，整个运算就失败。
    ∃i\exists i
  2. 冲突条件a_i ≠ b_i 且 a_i ≠ 1 且 b_i ≠ 1，两个维度既不相等也无法扩展。
    ai≠bi∧ai≠1∧bi≠1a_i \neq b_i \land a_i \neq 1 \land b_i \neq 1
  3. 失败结论满足上述条件即触发 shape mismatch 错误。
    error\text{error}

AI 应用场景

NLP 变长句子批处理

阶段：数据预处理 + 模型部署

解决的问题：BERT 类模型要求固定长度输入，但真实文本长短不一。传统做法要么逐句推理（极慢），要么无脑 pad 到超长（浪费算力）。AI 方案采用 max_len padding + attention mask，配合 jax.jit 编译，让 batch 高效处理且注意力只落在真实 token 上。

具体案例：一个文本分类服务每天收到 1~200 词的评论，统一 pad 到 256 后，用 jnp.arange(L) < valid_len 生成 mask，在 loss 计算中忽略填充位。JIT 编译后单 GPU 推理吞吐量提升约 3 倍，长尾输入也无需重新编译。

强化学习变长轨迹采样

阶段：训练数据管理

解决的问题：强化学习智能体每轮交互的步数不固定（终止时刻不同），直接堆叠成 batch 会形状不匹配。用 padding + mask 把轨迹统一到最大步数，mask 在 GAE 优势估计中屏蔽越界步，既保证 batch 形状静态，又不会让填充步污染梯度。

具体案例：在 Pendulum 环境中采样 200 条轨迹，长度在 180~220 之间波动。pad 到 220 后用 lax.cond 处理终止状态，配合 static_argnums 固定 horizon 超参，训练吞吐相比逐条更新提升 2 倍以上。

科学计算网格分辨率切换

阶段：模型编译优化

解决的问题：物理模拟中网格分辨率（如 64³、128³、256³）是固定的几档，每次切换都触发 recompilation 很浪费。用 static_argnums 把分辨率参数静态化，让每种分辨率各编译一次并缓存，运行时按需命中，避免反复重编译的开销。

具体案例：CFD 求解器在 64³ 与 128³ 网格间切换，静态化分辨率参数后，切换耗时从“分钟级重编译”降到“毫秒级缓存命中”，并行模拟效率显著提升。

Python 演示

python29 行
import jaximport jax.numpy as jnp # 1. 静态形状: 编译期完全确定, 一次编译多次复用def fixed_sum(x):    return jnp.sum(x) fixed_jit = jax.jit(fixed_sum)x = jnp.ones((4, 4))print("静态形状输出:", fixed_jit(x)) # 2. 变长序列: 固定最大长度 + mask 忽略填充位def masked_mean(x, valid_len):    mask = jnp.arange(x.shape[0]) < valid_len    mask = mask[:, None]    return jnp.sum(jnp.where(mask, x, 0.0)) / jnp.maximum(valid_len, 1) masked_jit = jax.jit(masked_mean)y = jnp.arange(12.0).reshape(6, 2)print("掩码均值(前3行):", masked_jit(y, 3))print("掩码均值(前5行):", masked_jit(y, 5)) # 3. static_argnums: 长度参数静态化, 参与 Python 切片def splice_sum(length, x):    return jnp.sum(x[:length]) static_jit = jax.jit(splice_sum, static_argnums=(0,))print("静态参数(截取4):", static_jit(4, x))print("静态参数(截取2):", static_jit(2, x))
目的：演示 JAX 中静态形状编译复用、padding+mask 处理变长输入、以及 static_argnums 强制参数静态化三种典型场景。
  1. import jax.numpy as jnp导入 JAX 的 NumPy 风格接口，所有数组操作都基于它。
  2. fixed_jit = jax.jit(fixed_sum)编译固定形状求和函数，形状 (4,4) 在编译期确定。
  3. mask = jnp.arange(x.shape[0]) < valid_len构造布尔掩码：位置索引小于真实长度则为 True。
  4. mask = mask[:, None]增加一维，使掩码能广播到 (L, E) 形状。
  5. jnp.sum(jnp.where(mask, x, 0.0))用掩码把填充位清零后再求和，填充位不参与计算。
  6. jnp.maximum(valid_len, 1)防止除零，同时让归一化分母是真实长度。
  7. static_argnums=(0,)把首个参数 length 视为编译期常量，值变化会触发重新编译。
  8. x[:length]因为 length 是具体 Python int，这里执行的是 Python 切片而非动态切片。

常见误区

误区我在 NumPy 里可以随便 reshape、拼接数组，JAX 应该也一样，形状在运行时想变就变。

正解JAX 的数组形状在编译（trace）阶段就已冻结，是静态属性。运行时“形状变化”本质上是新的静态形状触发了重新编译，或者通过 padding 模拟。形状是后厨账本上“固定的桌数”，不是散客随到随改的数字。

为何常见：NumPy 即时执行，形状天然是运行时属性；JAX 引入 JIT 后编译边界让形状冻结，初学者容易忽略这一层。

误区既然动态形状影响性能，我就永远不用它，所有输入一律 pad 到最大长度。

正解过度 padding 会浪费显存和计算，尤其是长度差异大的 batch，可能 90% 是填充位。合理做法是分桶（bucketing）按长度分组、按需使用 static_argnums、用 mask 控制有效计算。散客高峰期也可以“分时段预约”，而不是把所有桌子都按最大规模摆满。

为何常见：看到文档强调静态形状优势，容易走极端，忽视了灵活性本身的价值。

误区把尽可能多的参数标记为 static，JIT 编译就会变得更聪明、更快。

正解static_argnums 把参数变成编译期常量，参数值一变就整体重编译。比如把每轮都在变的 batch size 静态化，会导致每轮都重编译，性能反而雪崩。稳妥做法是只把“稳定且影响形状/控制流”的参数静态化，就像固定桌数的佛跳墙可以提前备菜，但散客桌不建议这么干。

为何常见：静态参数能参与形状推导，看起来“更优化”，但没意识到变量参数触发重编译的代价。

误区报错说 shape mismatch，就是我某个数字写错了，随便改一下就行。

正解shape mismatch 一般是广播规则破坏或维度对齐错误，要看完整形状元组逐维核验。用 x.shape 打印中间结果、对照“相等或一者为 1”的规则，才能精确定位是哪一维冲突。瞎猜数字等于不核对人数清单就认定“上错菜”，效率极低。

为何常见：错误信息只显示两个形状，初学者容易只盯着数字改，而不去理解对齐规则。

误区Python 支持分支和循环，JIT 编译的函数里应该也能直接写 `if length > 32:` 这样的代码。

正解tracer 的值在编译期不可知，Python 控制流要求条件必须是具体值（concrete value）。动态条件需要用 jax.lax.cond 或 jax.lax.fori_loop，或者把条件参数放入 static_argnums。主厨要知道客人人数才能决定做几人份，否则只能两种规格都准备——这就是 lax.cond 的工作方式。

为何常见：Python 的动态特性让人默认 if 总是可用，没意识到 JIT 追踪（trace）与执行的边界。

面试题

Q1. JAX 中数组的形状本质是什么？

单选题 · easy

  • A.运行时可以任意变化的属性
  • B.编译时确定的静态属性← 正确答案
  • C.仅与数据的具体值有关
  • D.由 GPU 显存大小决定

答案：B。JAX 数组形状在 trace 阶段确定，是静态属性。运行时“变化”本质是重编译或 padding 模拟，就像后厨账本上的桌数是提前锁定的。

Q2. 静态形状对 XLA 的核心价值是什么？

单选题 · easy

  • A.减少开发者写的代码量
  • B.让编译器能在编译期做 kernel 选择与内存布局优化← 正确答案
  • C.让 Python 解释器运行更快
  • D.避免使用 GPU 设备

答案：B。静态形状使 XLA 在编译期掌握完整形状信息，从而搜索最优 kernel、规划内存布局。正如提前知道 10 桌，主厨才能批量备菜、优化火力。

Q3. 把长度不一的文本序列输入 BERT 模型，最常用的处理方式是？

单选题 · easy

  • A.直接让模型报错
  • B.只保留最长的一条
  • C.padding 到最大长度加上 mask 忽略填充← 正确答案
  • D.随机截断到相同长度

答案：C。padding 到 batch 内最大长度让形状静态化，mask 标记真实位置使 loss 只关注有效 token。随机截断会丢失信息，直接报错则无法批处理。

Q4. 关于 jax.jit(f, static_argnums=(1,)) 的说法，正确的是？

单选题 · mid

  • A.第 1 个参数必须在运行时保持不变
  • B.第 2 个参数作为编译期常量，值变化会触发重编译← 正确答案
  • C.第 1 个参数的形状会动态变化
  • D.所有参数都必须声明为静态

答案：B。static_argnums 的索引从 0 开始，(1,) 指第二个参数。该参数变为编译期常量，值变化触发重编译，像固定桌数的佛跳墙菜单。

Q5. 广播规则中，形状 (3, 1) 与 (1, 4) 广播后的结果形状是？

单选题 · mid

  • A.(3, 4)← 正确答案
  • B.(4, 3)
  • C.(1, 1)
  • D.报错

答案：A。维度从右往左对齐，1 与 4 取 4，3 与 1 取 3，得到 (3, 4)。类似一人份和四人份合并成 3×4 的大拼盘。

Q6. 在 jit 函数内用动态 tracer 值作为 Python if 条件，会引发什么？

单选题 · hard

  • A.shape mismatch 错误
  • B.ConcretizationTypeError← 正确答案
  • C.内存溢出错误
  • D.自动广播警告

答案：B。tracer 的值在编译期不可知，Python 控制流要求具体值，于是触发具体化错误。需用 lax.cond 或 static_argnums 处理。

Q7. 关于广播规则，下列说法正确的有？

多选题 · mid

  • A.维度从右往左对齐比较← 正确答案
  • B.相同维度直接通过← 正确答案
  • C.维度一个为 1 时按另一个维度扩展← 正确答案
  • D.所有维度都必须完全相等才能运算

答案：ABC。A、B、C 是广播三原则。D 错误，维度为 1 时可扩展，不必完全相等——就像不同分量的食材也能合并成一道菜。

Q8. 关于 static_argnums 的正确说法有？

多选题 · hard

  • A.静态参数值变化会触发重新编译← 正确答案
  • B.静态参数可以让形状推导更精确← 正确答案
  • C.静态参数应该设为常变的 batch size
  • D.静态参数必须是可哈希的 Python 值← 正确答案

答案：ABD。A、D 是基本语义；B 是用途——静态参数可参与形状推导；C 错误，常变参数会导致频繁重编译，应只对稳定且影响形状的参数使用。

Q9. 广播规则中，形状 (4, 3) 与 (3, 4) 进行加法运算会？

单选题 · mid

  • A.广播为 (4, 4)
  • B.广播为 (3, 3)
  • C.报 shape mismatch 错误← 正确答案
  • D.广播为 (12, 12)

答案：C。从右往左第一维 3 与 4 不等且均非 1，破坏广播规则。就像 4 人份和 3 人份的菜无法直接合并成标准大份。

Q10. 关于 JAX 的形状推导，下列说法正确的是？

单选题 · extreme

  • A.每个算子的输出形状都能由输入形状纯函数地推导← 正确答案
  • B.形状推导只在运行时根据数据值进行
  • C.动态形状可以跨设备自由传输而无需重编译
  • D.形状推导结果取决于数据的数值内容

答案：A。JAX 每个算子都有形状推导规则，是 shape → shape 的纯函数，不依赖具体数值。B、D 错误；C 中动态形状变动仍需重编译。