深夜的魔法菜谱馆里，架子上摆满封面编着号码的菜谱。每一页都印着一串看似毫无规律的编号，客人以为这是厨师随手抓的，其实早在印刷时，这些编号就被一台“搅乱机”按固定规则排好了：同一本菜谱，无论你翻多少次，第 10 页永远是那个编号，第 11 页也永远是下一个。真正的随机像掷骰子，每次都不一样；而菜谱的编号是“看似随机、其实固定”。这套机制，就是计算机里的伪随机数生成器（PRNG，Pseudo-Random Number Generator）。 老式厨房只有一本共享菜谱挂在灶台边，所有厨师翻来翻去，谁先翻、谁后翻，完全说不清——这是传统“全局随机种子”的混乱。JAX 的解法很特别：扔掉共享菜谱，每位厨师手握自己的魔法菜谱（key）。要用随机数时，先让菜谱“分裂”出新的菜谱（split），旧菜谱继续留在自己手里。因为每本菜谱的内容由封面编号唯一决定，所以只要记住初始编号（seed），整场宴席的每一个数字都能精确重现。这篇文章，就带你走进这座菜谱馆，从窗口看到后厨，把 PRNG 的整套机制一次看懂。

岗位关联

学会 JAX 的 PRNG 设计，你就能写出可复现、可并行、可调试的深度学习实验，避免“随机数失控”这类隐性 Bug。这是从“代码能跑”走向“实验可靠”的关键一步，也是面试与工程评审中的高频考点。

目录
  1. 可视化演示
  2. PRNG基础概念：魔法菜谱与确定性随机
  3. 全局种子困境：共享菜谱的混乱
  4. 显式状态设计：把状态握在自己手里
  5. key与seed：密钥的构造
  6. split分裂机制：随机数流的分叉
  7. Threefry原理：搅乱机内部
  8. 工程实践要点：可复现的炼丹流水线
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

展示JAX显式PRNG状态的密钥分裂机制，以及显式密钥如何保证随机数序列的可复现性和并行安全。

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

PRNG密钥分裂树 种子生成初始PRNGKey，内部存储两个无符号整数 split函数将父密钥分裂为多个统计独立的子密钥 每个子密钥生成独立的随机数序列，互不干扰 split(key, num) → [key_0, key_1, ..., key_{n-1}] 显式密钥传递确保函数纯性。

PRNG基础概念：魔法菜谱与确定性随机

想象一座深夜营业的魔法菜谱馆。架子上摆满了封面编着号的菜谱，每一页都印着一串看似毫无规律的编号。客人以为这些编号是厨师随手抓的，其实它们早在印刷时就被一台“搅乱机”按固定规则排好了：同一本菜谱，无论你翻多少次，第 10 页永远是那个编号。真正的随机（比如掷骰子）每次都不一样，而菜谱编号是“看似随机、其实固定”。这套机制，就是计算机里的伪随机数生成器 PRNG。

PRNG 不产生真正的随机，它只从一个初始状态出发，按确定性的数学规则不断产出“看起来随机”的数列。这个初始状态就是随机种子 seed，也就是你翻开菜谱的那一页。JAX 底层负责“搅乱”的机器叫 Threefry，但我们先不急着拆机器，先把“种子 + 递推规则”这个框架看清楚。好的 PRNG 要能通过统计检验：输出在均匀性、独立性上和真正的随机数难分彼此，才能放心用在 Dropout、初始化、数据打乱这些场景里。

最经典的伪随机发生器是线性同余法（LCG）：拿到上一个数后，依次做乘、加、取模三步，就得到下一个数。只要三个常数选得讲究，输出的序列在统计上就和均匀分布几乎无法区分。

菜谱的价值恰恰在于可重现：深度学习里，数据划分、参数初始化、Dropout 都需要“这次的随机和上次一样”。伪随机不是骗人，而是把“随机”变成一种可回放的资产——这正是 PRNG 在 AI 工程里的地基意义。

xn+1=(a×xn+c) mod m\displaystyle x_{n+1} = (a \times x_n + c) \bmod m
通俗：上一个数字经过“乘、加、取模”三步变成下一个数字。只要不告诉别人初始值，这串数字看起来就和抛硬币一样乱。
  1. 乘法当前值乘以常数 a，把相邻两项的差异快速拉开，避免序列挤在一起。
    xn×ax_n \times a
  2. 加法再加上常数 c，防止序列退化成等差数列或全零这类病态模式。
    xn×a+cx_n \times a + c
  3. 取余对模数 m 取余，把结果压回 [0, m) 区间，保证数值不溢出且周期可控。
    (xn×a+c) mod m(x_n \times a + c) \bmod m

全局种子困境：共享菜谱的混乱

老式厨房只有一本共享菜谱，挂在灶台边任人翻动。王师傅翻到第 10 页取了编号，李师傅接着翻到第 11 页，张师傅又翻回第 10 页……等王师傅下次想要“第 10 页的下一道菜”时，页面早被翻乱了，他拿到的编号和上次对不上。整个厨房的随机序列取决于谁先翻、谁后翻，完全不可预测，更不可复现。

numpy、Python 标准库 random 正是这样的“共享菜谱”：一个进程内只有一个全局随机状态，任何代码调用 np.random.rand() 都会偷偷改动它。深度学习里，数据加载、参数初始化、Dropout、数据增强到处都在抽随机数，彼此没有任何隔离。更隐蔽的是批处理污染：哪怕只是多跑了一个 batch 的数据增强，后续所有 batch 的采样顺序都会被改写，整个实验的轨迹随之漂移。

理想的可复现条件是：第 i 个样本只取决于种子和序号 i，与“之前谁调过随机”无关。但全局状态下，所有历史调用都会影响下一个输出：

共享菜谱让并行彻底瘫痪：8 个厨师同时翻同一本菜谱，谁先谁后根本说不清，复现实验变成了买彩票。于是 JAX 决定：扔掉公共菜谱，每位厨师自带一本——这就是显式 PRNG 状态的起点。

理想: xi=g(seed,i)全局状态: xi=g(seed,i,ht)\displaystyle \text{理想: } x_i = g(\text{seed}, i) \qquad \text{全局状态: } x_i = g(\text{seed}, i, h_t)
通俗：h_t 是“到目前为止被翻过的全部页数记录”。别人多翻了一页，你拿到的数字就变了。
  1. 理想依赖样本只由种子与序号 i 决定，与历史无关，因此可复现。
    xi=g(seed,i)x_i = g(\text{seed}, i)
  2. 现实依赖全局状态让输出多了一个历史变量 h_t，即此前所有调用的副作用记录。
    xi=g(seed,i,ht)x_i = g(\text{seed}, i, h_t)
  3. 破坏后果h_t 由调用顺序决定，任何代码改动都会改变它，复现因此失败。
    ht≠hs⇒xi 前后不一致h_t \neq h_s \Rightarrow x_i \text{ 前后不一致}

显式状态设计：把状态握在自己手里

JAX 的厨房改革很简单：做菜必须自己带菜谱，菜谱是“入参”也是“出参”。厨师把菜谱递给采样窗口前，窗口会基于这本菜谱产出随机样本；想要后续的随机数，就先用 split 派生一本新菜谱。旧菜谱严格封装，绝不允许别人偷偷翻页。因为没有共享菜谱，任何两位厨师之间天然隔离。

这对应 JAX 的“函数式纯随机”设计：随机状态（key）被显式地作为参数传入、传出，而不是藏在全局变量里。随机函数是纯函数——同样的 key 进去，必然产生同样的样本，绝无隐藏副作用。传统框架里“调用一次随机数，全局状态推进一格”的隐式逻辑，在 JAX 里被彻底取消了。

纯函数性质可以写成一句话：输出完全由输入的 key 决定。

每位厨师只对自己的菜谱负责，菜谱不怕被复制、不怕被并行。深度学习训练里的数据并行、流水并行，从此可以放心地把不同菜谱分给不同设备，互不干扰——这是 JAX 在分布式训练中广受欢迎的重要原因。

k1=k2⇒x1=x2;k1≠k2⇒x1⊥x2\displaystyle k_1 = k_2 \Rightarrow \mathbf{x}_1 = \mathbf{x}_2;\qquad k_1 \neq k_2 \Rightarrow \mathbf{x}_1 \perp \mathbf{x}_2
通俗：同一本菜谱读两遍，菜完全一样；两本不同的菜谱，做出来的菜互不相干。
  1. 同一输入相同 key 进入随机函数，函数内部不读取任何外部状态。
    k1=k2k_1 = k_2
  2. 同一输出得到的样本必然完全相同，这是纯函数的核心承诺。
    x1=x2\mathbf{x}_1 = \mathbf{x}_2
  3. 不同密钥独立不同 key 产生的样本统计上互不相关，可放心分发。
    k1≠k2⇒x1⊥x2k_1 \neq k_2 \Rightarrow \mathbf{x}_1 \perp \mathbf{x}_2

key与seed：密钥的构造

新厨房里“菜谱”的正式名称叫 key（密钥）。它像一本封面刻着长长编号的魔法菜谱，编号决定了里面每一页的内容。你没法直接“打印一本随机内容的菜谱”，但你可以指定一个初始页码 seed，由印刷机据此生成一本唯一的菜谱。

JAX 中构造密钥的入口是 jax.random.key(seed)。seed 可以是你随手写的整数，比如 42；返回的是一个形状为 (2,) 的 uint32 数组，这就是 key。注意 key 本身不是随机数，它是“随机数的来源”——就像菜谱是菜的来源，而不是菜。key 一旦生成就应视为不可变对象，不要原地修改。

key 是 64 位空间（两个 32 位整数拼成），不同 seed 撞出同一个 key 的概率微乎其微：

第一次进入菜谱馆，你只需记住初始封面编号 seed=42，此后无论谁、无论在哪台机器上，只要用同一编号印刷菜谱，整场实验的随机旅程就完全一致。另外，key 也支持批量形式：传入形状为 (batch, 2) 的密钥组，可以一次为整批样本生成独立抽样，配合 vmap 使用非常顺滑。

k=key(seed);Pr⁡[k(sa)=k(sb)]≈2−64 (sa≠sb)\displaystyle k = \mathrm{key}(\text{seed});\qquad \Pr\big[k(s_a) = k(s_b)\big] \approx 2^{-64} \ (s_a \neq s_b)
通俗：世界上几乎不可能有两本封面编号完全相同的菜谱。
  1. 构造密钥把整数种子映射为固定形状的密钥数组。
    k=key(42)k = \mathrm{key}(42)
  2. 密钥结构key 是长度为 2 的 uint32 数组，本质是两段 32 位整数。
    k∈Z2322k \in \mathbb{Z}_{2^{32}}^2
  3. 碰撞概率不同种子撞出同一密钥的概率约为 2^{-64}。
    Pr⁡[k(sa)=k(sb)]≈2−64\Pr[k(s_a) = k(s_b)] \approx 2^{-64}

split分裂机制：随机数流的分叉

单本菜谱只支持“顺序翻页”，但深度学习需要几百个并行厨师同时做菜。魔法菜谱的妙处在于：当你需要更多厨师时，可以把现有菜谱“复印”成两本新菜谱——但复印机做了手脚，两本新菜谱的封面编号被重新打乱，内容互不相同、也与原版无关。原版照旧留着继续用，谁也不用等谁。

这就是 jax.random.split(key, num)：它把一把 key 变成 num 把新的 key，使用新 key 产生的随机流彼此独立。典型用法是“一留一用”：

split 可以反复进行，形成一棵“密钥树”：根 key 分裂出 2 把，各再分裂出 2 把，得到 4 把……第 d 层有：

数据并行时，主进程把菜谱层层分发：每个 worker 拿到一片独一无二的叶子菜谱，做自己的数据增强、Dropout，互不串味。只要根菜谱的编号不变，整棵树的每一片叶子都完全可复现——这就是“随机数流控制”的核心。

N(d)=2d;树形层级: k→{k1,k2}→{k11,k12,k21,k22}\displaystyle N(d) = 2^{d};\qquad \text{树形层级: } k \to \{k_1, k_2\} \to \{k_{11}, k_{12}, k_{21}, k_{22}\}
通俗：一本菜谱可以生出两本新菜谱，再生出四本……像细胞分裂一样指数增长，每片叶子都独立做菜。
  1. 分裂操作一把 key 生成两把新 key，父 key 本身不受影响。
    k↦(k1,k2)k \mapsto (k_1, k_2)
  2. 独立性质子密钥各自的随机流统计上独立。
    flow(k1)⊥flow(k2)\mathrm{flow}(k_1) \perp \mathrm{flow}(k_2)
  3. 指数增长分裂 d 层后叶子密钥数为 2^d。
    N(d)=2dN(d) = 2^d
python3 行
key = jax.random.key(42)key, subkey = jax.random.split(key)      # 派生一把子钥x = jax.random.normal(subkey, (3,))

Threefry原理：搅乱机内部

封印在菜谱封面编号里的“搅乱机”长什么样？它是一台由齿轮和拨片构成的机械：你把两个数字从左边投进去——封面编号 K 和当前页计数器 C，摇动若干圈，右边就会掉出一个面目全非的数字。搅乱机内部是固定的齿轮组合，同样的输入永远得到同样的输出；但你无法从输出反推输入，因为齿轮里充满异或和循环位移。

JAX 底层默认使用 Threefry 算法，它属于“计数器型”伪随机生成器。核心思路：把 key 当作密钥，把 counter 当作输入，用一个可逆置换把 (key, counter) 映射成看似随机的输出。为了让输出与 counter 的关系足够复杂，最后还要把置换结果与 counter 本身做异或。

单轮混合操作可以简写为对两个状态字的操作：

整个 Threefry 输出为：

这台搅乱机的意义在于“确定性加不可预测感”：菜谱每翻一页，计数器加一，搅出的编号看似随机却由封面编号唯一决定。正因核心是固定机械结构，JAX 才能保证同一 key、同一 counter 在任何设备上都输出完全一致的随机序列。

x+=y;y=rotl(y,r)⊕x\displaystyle x \mathrel{+}= y;\qquad y = \mathrm{rotl}(y, r) \oplus x
O=C⊕Threefry(K,C)\displaystyle O = C \oplus \mathrm{Threefry}(K, C)
通俗：先把（计数器，密钥）塞进搅乱机来回揉，再把揉出来的结果和计数器对拧一次，让输出与计数器之间看不出线性痕迹。
  1. 字混合两个字先相加，其中一个再循环左移 r 位并异或另一个，实现单次搅动。
    x+=y;y=rotl(y,r)⊕xx \mathrel{+}= y;\quad y = \mathrm{rotl}(y, r) \oplus x
  2. 多轮迭代上述搅动重复若干轮（如 20 轮），每四轮注入一次与密钥相关的轮常数。
    Or+1=R(K,Or)O_{r+1} = R(K, O_r)
  3. 异或回收把最终置换结果与原始计数器异或，破坏可逆线索。
    O=C⊕Threefry(K,C)O = C \oplus \mathrm{Threefry}(K, C)

工程实践要点：可复现的炼丹流水线

宴会办得越大，越需要一个“总管家”：根菜谱只由主厨保管，编好号。每个部门（数据、模型、评估）要用随机时，不能自己去复印，而是向总管家申领一个“域编号”，由总管家用 fold_in 机制把根菜谱和域编号结合，生成专属子菜谱。这样一来，任何部门的随机流都可以单独定位、单独复现。

JAX 提供 jax.random.fold_in(key, data)：它把根 key 与一个整数 data 绑在一起派生新 key，专用于“同一根密钥下按 ID 派生独立子流”。相比反复 split，fold_in 更适合模块化：只要根 key 与模块 ID 不变，模块内的随机行为就稳定，适合按数据模块、模型模块、评估模块划分领地。

工程上推荐“一处入口、处处派生”：

复现实验时只需记录根种子，所有模块的子流自动恢复。再配合版本锁存——记录 JAX 与硬件版本，即使将来 API 变化也能定位差异。

总管家模式还让“断了也能续”：训练中断时把当前 key 存进 checkpoint，恢复时从该 key 继续翻菜谱，后续随机序列与不中断时完全一致。这把 PRNG 从“玄学”变成了工程上可审计、可回放的基础设施。

kmodule=fold_in(kroot,module_id)\displaystyle k_{\mathrm{module}} = \mathrm{fold\_in}(k_{\mathrm{root}}, \mathrm{module\_id})
通俗：记录根菜谱编号加各部门编号清单，整个宴会的菜序就能原样重演。
  1. 根密钥整个实验只从一个根种子构造全局唯一密钥。
    kroot=key(base_seed)k_{\mathrm{root}} = \mathrm{key}(\mathrm{base\_seed})
  2. 模块派生fold_in 把模块 ID 与根密钥绑定，生成稳定可复现的子流。
    kmodule=fold_in(kroot,module_id)k_{\mathrm{module}} = \mathrm{fold\_in}(k_{\mathrm{root}}, \mathrm{module\_id})
  3. 复现闭环记录 base_seed 与全部 module_id，整场实验即可原样重演。
    run(kroot)≡run′(kroot)\mathrm{run}(k_{\mathrm{root}}) \equiv \mathrm{run}'(k_{\mathrm{root}})

AI 应用场景

多卡数据并行训练

阶段：模型训练（分布式训练）

解决的问题：传统全局随机种子在 8 卡并行时无法隔离各卡的随机流——谁先调用随机函数都会污染他人，复现实验全靠运气。JAX 用 split 把根密钥分裂成多把独立子钥，每张卡持有一把，随机流互不干扰且整体可复现。

具体案例：在 8×A100 上训练大语言模型时，主进程执行 jax.random.split(root, 8) 得到 8 把 worker 子钥，分别用于各自的 Dropout、数据增强与参数初始化；实验报告只记录根种子，任何一次训练都能精确复原。

大规模超参搜索的实验对比

阶段：实验配置与调参（AutoML / 超参搜索）

解决的问题：并发跑几十组超参时，如果用共享全局种子，各组的数据划分与初始化会互相串扰，导致对比不公平。用 fold_in(root, trial_id) 为每个 trial 派生独立随机流，保证组间公平、单个 trial 可复现。

具体案例：算法团队用 Optuna 同时发起 100 组学习率与 batch size 搜索，每组用 fold_in 派生独立的数据划分与权重初始化；负责人只需保存根种子，就能随时回放任意一组的完整随机过程。

强化学习训练的可复现回放

阶段：数据生成与环境仿真（强化学习）

解决的问题：RL 训练中环境随机性与策略随机性必须分离——环境种子变了，同一策略的回报就无法比较；策略种子变了，同一环境下的探索行为也无法复现。用 split 把根密钥分成 env_key 与 policy_key 两棵流，互不污染。

具体案例：训练 CartPole 时把 env_key 用于环境初始状态与扰动，policy_key 用于动作采样；中断后保存当前 key，恢复训练即可从断点继续生成与未中断完全一致的轨迹序列。

Python 演示

python39 行
import jaximport jax.numpy as jnp  def main():    # 根种子统一入口    root = jax.random.key(42)     # 按模块 ID 派生独立子流    data_key = jax.random.fold_in(root, 0)    model_key = jax.random.fold_in(root, 1)    eval_key = jax.random.fold_in(root, 2)     # 数据模块: 打乱并划分训练测试集    x = jnp.arange(1000)    perm = jax.random.permutation(data_key, x)    train, test = perm[:800], perm[800:]    print("train size:", train.shape[0], "test size:", test.shape[0])     # 模型模块: 小随机数初始化权重    w = jax.random.normal(model_key, (4, 4)) * 0.02    print("init weight mean:", float(w.mean()))     # 评估模块: 对测试集加噪声模拟扰动    noise = jax.random.normal(eval_key, (200,))    print("eval noise std:", float(noise.std()))     # split 派生出 8 个并行 worker 的子钥    keys = jax.random.split(data_key, 8)    print("num worker keys:", len(keys))     # 复现验证: 相同根种子再次运行结果一致    root2 = jax.random.key(42)    w2 = jax.random.normal(jax.random.fold_in(root2, 1), (4, 4)) * 0.02    print("reproducible:", bool(jnp.allclose(w, w2)))  if __name__ == "__main__":    main()
目的：演示 JAX 的 PRNG 工程范式——根种子统管、fold_in 按模块派生、split 支持并行、同种子精确复现。
  1. root = jax.random.key(42)全实验唯一根密钥，只此一处入口。
  2. data_key = jax.random.fold_in(root, 0)用模块 ID 0 派生数据专属子流。
  3. perm = jax.random.permutation(data_key, x)数据流负责打乱，保证划分确定性。
  4. w = jax.random.normal(model_key, (4, 4)) * 0.02模型流负责初始化，与数据流隔离。
  5. keys = jax.random.split(data_key, 8)一次分裂出 8 把 worker 子钥。
  6. w2 = jax.random.normal(jax.random.fold_in(root2, 1), (4, 4)) * 0.02重新构造根种子后结果完全一致。

常见误区

误区生成一个 key 打印出来是一串数字，那这串数字不就是随机数吗？直接拿它当随机样本用不就行了？

正解key 是“随机数的来源”，不是样本本身。就像菜谱是菜的来源而不是菜。key 必须喂给 normal、uniform、permutation 等采样函数才会产出样本；直接打印 key 看到的只是一段密钥数组，其数值分布与目标分布毫无关系。更关键的是，同一把 key 重复喂给采样函数，得到的是完全相同的样本。

为何常见：传统 API 叫 random.seed 且不返回值，新手自然把“种子”和“随机数”混为一谈；JAX 把种子显式变成对象后，这种混淆更容易发生。

误区我用同一个 key 调用两次 `normal`，第二次应该得到不同的随机数吧？毕竟函数名叫 random。

正解纯函数特性决定了同一 key 必然产生同一输出。要得到新的随机样本，必须先用 split 派生新 key，或者循环里维护 key 的更新。把同一把 key 用两次，等于把同一页菜谱读两遍，菜完全相同——这不是 bug，而是 JAX 的设计核心。

为何常见：习惯了 numpy 全局状态“每次调用都自动推进”的语义，忽略了 JAX 无状态、纯函数的设计。

误区split 就是把一把 key 复制成两份吧？那两份 key 产生的随机数应该一模一样。

正解split 是“派生且打乱”：子 key 与父 key、子 key 与子 key 之间产生的随机流彼此独立（近似独立）。比喻中的复印机做了手脚，两本新菜谱被重新装订。这正是 JAX 能安全地把 key 分发给不同 worker 的根本原因——如果只是复制，并行分发就毫无意义。

为何常见：“分裂”这个词容易让人联想到“复制切片”，而忽略了内部对密钥进行了重新混合。

误区我在训练脚本开头调了 `np.random.seed(0)`，又调了 `random.seed(0)`，JAX 的随机结果应该也能复现吧？

正解JAX 的随机完全不受 numpy 和 Python 标准库 random 的影响，必须用 jax.random.key(seed) 构造根密钥并显式传递。全局种子只能控制传统库，两者是两套独立系统。比喻：你锁好了旧厨房的共享菜谱，但新厨房的每位厨师根本不看那本。

为何常见：深度学习各框架的随机机制不同，从 PyTorch 或 TensorFlow 迁移习惯时容易张冠李戴。

面试题

Q1. JAX 中构造根密钥的正确 API 是？

单选题 · easy

  • A.jax.random.seed(42)
  • B.jax.random.key(42)← 正确答案
  • C.jax.random.normal(42)
  • D.np.random.seed(42)

答案：B。jax.random.key(seed) 是 JAX 构造显式密钥的标准入口。A 和 D 都是传统全局种子体系，与 JAX 的显式 key 设计无关；normal 是采样函数，需要先传入 key。比喻：只有 B 是给新厨房印刷专属菜谱，其余都是旧厨房的锁。

Q2. 在 JAX 中，用同一把 key 调用两次 jax.random.normal，会得到？

单选题 · easy

  • A.两组完全不同的随机数
  • B.第一组与第二组部分相同
  • C.两组结果完全相同← 正确答案
  • D.程序会直接报错

答案：C。JAX 的随机函数是纯函数，相同输入必然相同输出，所以两次结果完全相同。要想得到不同的随机样本，必须用 split 派生新 key。比喻：同一页菜谱读两遍，菜完全一样。

Q3. 对一把根 key 连续执行 split 三次（每次都把所有子钥继续分裂），一共能得到多少把叶子密钥？

单选题 · mid

  • A.3 把
  • B.6 把
  • C.8 把← 正确答案
  • D.16 把

答案：C。每次 split 使 key 数量翻倍：1 → 2 → 4 → 8，三次后得到 2^3 = 8 把叶子密钥。这对应密钥树的层级结构，也解释了为什么 split 能高效支撑大规模并行。

Q4. JAX 随机数与 numpy 随机数的本质区别是？

单选题 · mid

  • A.JAX 速度更快所以更随机
  • B.JAX 把随机状态显式作为参数传递← 正确答案
  • C.JAX 使用真正的硬件随机源
  • D.numpy 无法产生可复现结果

答案：B。JAX 采用函数式显式状态设计，key 作为纯函数入参，无全局副作用；numpy 依赖全局隐式状态。JAX 的随机同样是伪随机、同样可以复现，并不依赖硬件随机源，速度与“随机性”无关。

Q5. 需要为 8 个数据并行 worker 分配互不干扰的随机流，最合适的做法？

单选题 · mid

  • A.各 worker 都调 np.random.seed(0)
  • B.在根 key 上 split 出八把子钥← 正确答案
  • C.每个 worker 使用同一把 key
  • D.直接用固定整数 8 当随机数

答案：B。用 root_key = jax.random.split(root, 8) 一次产出 8 把独立子钥，分发给 8 个 worker。C 会让所有 worker 得到相同样本，A 是全局种子不可控，D 根本不是随机数。

Q6. Threefry 算法中，为了防止输出与计数器存在线性关系，最后一步通常做什么？

单选题 · hard

  • A.把输出取平均降噪
  • B.把置换结果与计数器异或← 正确答案
  • C.把计数器清零重新计数
  • D.对输出做 softmax 归一化

答案：B。Threefry 将可逆置换结果与原始 counter 异或，即 O = C ⊕ Threefry(K, C)，破坏输出与计数器的线性结构，让输出看起来足够随机。比喻：搅乱机最后再对拧一次齿轮，让痕迹彻底消失。

Q7. fold_in(root, id) 与 split(root) 的主要分工差异是？

单选题 · hard

  • A.fold_in 更快但结果不可复现
  • B.fold_in 适合按固定 ID 派生可复现子流← 正确答案
  • C.split 只能使用一次不能反复分裂
  • D.两者功能完全相同可以互换

答案：B。fold_in 将根 key 与固定整数 ID 绑定，生成稳定可复现的子流，适合模块化管理（数据、模型、评估各自分域）；split 面向并行分叉与层级树。两者都能复现，split 也可反复使用。

Q8. 关于 JAX 中 key 的使用，下列哪些说法正确？

多选题 · mid

  • A.key 可以重复使用得到相同样本← 正确答案
  • B.key 用完后应通过 split 更新← 正确答案
  • C.key 可以直接当作随机样本
  • D.同一 key 用于不同采样函数← 正确答案

答案：ABD。纯函数下同一 key 必得同一输出，A 说法正确但实践中应通过 split 更新 key，所以 B 也对；key 是随机数来源而非样本，C 错；同一 key 喂给 normal 与 uniform 可得到不同分布的样本，D 对。

Q9. 需要完整复现一次分布式训练实验，以下哪些做法是必要的？

多选题 · hard

  • A.统一由根种子派生子流← 正确答案
  • B.记录 JAX 框架版本号← 正确答案
  • C.使用全局 np.random.seed 控制全流程
  • D.保存模块的 fold_in ID← 正确答案

答案：ABD。复现依赖三要素：根种子统一入口、各模块派生 ID 记录、框架版本锁存。np.random 的全局种子与 JAX 无关，无法控制 JAX 的随机行为，因此 C 不必要。

Q10. 关于 Threefry 的工程特性，下列哪些理解正确？

多选题 · extreme

  • A.计数器型 PRNG，输出由 key 与 counter 决定← 正确答案
  • B.同 key 同 counter 在不同设备输出一致← 正确答案
  • C.Threefry 可逆因此可以直接反推随机序列
  • D.多轮混合用加法旋转异或实现雪崩← 正确答案

答案：ABD。Threefry 是计数器型、确定性的，输出由 key 与 counter 共同决定，且设备无关（A、B 对）。虽然内部置换在数学上可逆，但经过多轮打散与最终异或后，实用上无法从输出反推序列（C 错）。轮内混合使用加法、循环左移与异或，实现雪崩效应（D 对）。