写 JAX 代码的人迟早会被随机数生成上一课。我当初从 numpy.random 迁移过来第一行jax.random.normal(0.0, 1.0)就直接报错骂了半天的同时也第一次意识到JAX 的随机数不是换个库名那么简单而是一次从过程式到函数式范式的彻底重构。这背后牵涉的是显式状态传递、counter-based 算法和一套关于确定性的全新约定——也就是标题里那场确定性革命的真正含义。无论你是用 JAX 跑深度学习、做贝叶斯采样还是像最近很火的 whisper-jax 那样把模型打包成大规模批处理服务理解这套随机数体系都会帮你少踩无数个静默的坑真正掌控实验的可复现性。这篇不是文档复读我会尽量把为什么这样设计实际怎么用坑在哪讲透。适合两类人刚上手 JAX 一直被随机 API 搞晕的新手以及已经在用、但偶尔被同一个 key 怎么又生成了同样的数据折磨的进阶用户。先把结论摆出来JAX 的随机数和 numpy 的随机数差的不只是函数名而是整套关于随机状态由谁保管的哲学。1. numpy.random 的全局状态到底坑在哪1.1 一段代码里的隐形状态Mersenne Twister 的账本模型numpy.random 的底层是 Mersenne Twister它维护一个长度为 624 的 uint32 内部状态数组每次调用np.random.rand()都会读取并推进这个全局状态。用账本来类比最贴切整个进程共享一本账谁调用谁就在同一行往下记账。np.random.seed(42)只是把账本擦干净重写第一行之后任何模块、任何库的调用都会继续往下写。这就带来两个特别隐蔽的问题。第一个是调用顺序即命运——你导入某个库的顺序、某个工具函数是否内部用了np.random都会改变你后面所有随机结果。我遇到过数据增强库在后台偷偷调用np.random.randint导致同样的 seed 在不同环境里训练结果完全不同。排查这种问题极其痛苦因为报错信息为零只有把调用链彻底理清才发现是隔壁邻居污染了你的随机流。第二个问题是复制状态的伪独立。np.random.get_state()可以拿到这份全局状态很多并行框架为了复现会把状态拷贝给子进程结果每个子进程拿到的是完全相同的副本产出的随机序列逐位重合。我见过一个多进程数据加载代码四个 worker 采样出来的 batch 几乎一模一样训练指标还一路在涨——因为模型确实在学会那四份重复的数据只是泛化能力一塌糊涂。这种错误极难定位单看每个子进程一切正常对比起来才发现所有 worker 都在消费同一段随机流。1.2 为什么 JAX 必须推翻重来要理解 JAX 的随机数设计先得理解 JAX 的底层逻辑它假设所有函数都是纯函数——相同输入永远产生相同输出没有任何隐藏副作用。这个假设是 XLA 编译、自动微分、向量化vmap和多设备并行pmap的地基。一个带隐式全局状态的随机函数在编译器眼里是不可预测的黑洞没法被安全缓存没法被并行复制也没法在 GPU 上执行因为状态在 CPU 内存里每次调用都要同步。PyTorch 怎么处理它提供了torch.Generator这个显式对象你可以传入 generator 控制随机流如果你不给它依然退回全局状态。TensorFlow 也有类似的历史包袱。JAX 的做法更彻底根本没有全局种子随机函数必须显式接收一个 key 作为第一个参数。随机性不再是环境属性而是函数签名里的一个普通参数。这一刀切得干净代价是刚上手不习惯——但一旦理解你会发现并发、复现、向量化原本纠缠不清的问题被一次性拆掉了。2. JAX 随机数生成的核心设计Key、split 与底层算法2.1 PRNGKey 不是种子而是一把加密密钥jax.random.key(42)返回的不是普通种子而是一个 typed key在默认配置下本质是uint32[2]数组。它代表什么在 JAX 使用的 counter-based PRNG计数器型伪随机数生成器里这个 key 实际是加密算法的密钥。counter-based 的思想很直接拿一把密钥去加密一个不断递增的计数器加密结果的 bit 就是随机数。密钥不同密文完全不同密钥相同密文完全一致。第一步要扭转的观念是JAX 不靠推进状态来产生下一个随机数。Mersenne Twister 让状态不断前进同一个随机流是一条线counter-based 的世界里一把 key 本身就拥有一个巨大的独立随机空间——Threefry 的计数器空间是 2^64 个 block 量级每块产出 128 bit你在任何现实场景中都消费不完。你说random.normal(key, shape(1000,))它直接在这把 key 对应的独立流里取 1000 个样本根本不需要 split。所以key 的概念更接近身份证明而不是当前位置。这也是官方文档反复强调的同一个 key 在同一个算法下生成的随机数完全确定你可以任意重新计算结果逐位一致。确定性在 JAX 里不是限制恰恰是特性。2.2 split 与 fold_in独立流是派生的不是移动的既然一把 key 就有海量独立流那 split 是干嘛的答案是创建互不相关的多把 key。random.split(key, n)把一把 key 变成 n 把新的 key它们两两统计独立后续从任何一把子 key 生成的随机序列都不会与其他子 key 的序列重叠或相关。这正好对应并行场景每个样本、每个设备、每个训练步需要的是独立噪声而不是同一个流的切片。另一个常用函数是random.fold_in(key, data)它根据一个整数 data 确定性地产出一把新 key。它和 split 的区别在于fold_in 是可寻址的——给定父 key 和数据标识任何人都能重新派生同一把子 key适合按样本 ID、batch 号、epoch 号寻址的数据增强。而 split 产生的是匿名子 key事后无法重建。用生活类比split 像给每张身份证办独立户口fold_in 像用门牌号反查地址——两者都是派生态但可寻址性完全不同。2.3 三种底层生成器Threefry、Philox 与 rngJAX 内置了多种可切换的底层算法接口完全不变只是底层随机序列不同。主流讨论集中在三种生成器算法来源特点适用场景Threefry默认Threefish 分组密码Skein 哈希统计质量极高通过 TestU01 BigCrush理论安全性强默认通用统计实验、科学计算PhiloxRandom123 系列计数器型同样是加密级质量GPU 上通常更快GPU/TPU 大规模并行性能敏感rngXorshift 变体速度最快质量满足一般 ML 需求对速度极敏感、统计要求不极端的场景切换方式有两种全局配置jax.config.update(jax_default_prng_impl, philox)或者生成 key 时指定random.key(seed, impl...)。注意不同实现的数值序列完全不兼容——你拿着 Threefry 生成的 key 去 Philox 下复现结果是另一套。所以不要在生产环境随意切换 impl否则之前保存的实验结果全部失去对照意义。默认 Threefry 是 JAX 团队反复权衡后的稳妥选择绝大多数人根本不需要换。2.4 为什么同一个 key 调用两次结果是完全一样的这是新手最容易踩的坑也是函数式设计最集中的体现。写下a random.normal(key, shape(3,)) b random.normal(key, shape(3,))你会发现a和b逐位相等。在 numpy 里连续调用两次np.random.normal()会得到不同结果因为全局状态前进了在 JAX 里random.normal是纯函数输入没变输出必然不变。想要第二组不同样本必须先把 key split 一次key, subkey random.split(key) a random.normal(key, shape(3,)) b random.normal(subkey, shape(3,))这个设计带来的隐藏好处是你可以安全地重新计算任意中间结果。调试时想重跑某一步不用担心随机流被消费掉了。坏处是不养成 split 习惯你会生产大量重复的随机数据。后面我会专门讲怎么把这个习惯内化成肌肉记忆。提示如果你发现自己写了第二处random.normal(key, ...)而中间没有 split几乎可以断定这是一个随机流复用 bug。3. jax.random 实操从入门到批量并行3.1 最小可用示例生成第一个正态分布先看最基本的代码import jax import jax.numpy as jnp from jax import random # 新 API推荐 key random.key(42) # 旧 API已弃用但兼容 # key random.PRNGKey(42) samples random.normal(key, shape(5,)) print(samples)注意三点。第一key 是第一个位置参数后面才跟分布参数和 shape这与 numpy 的函数签名完全不同。第二random.key返回带类型的 typed key而random.PRNGKey返回裸的uint32[2]数组两者在大多数场景可以互换但新代码一律用random.key避免将来版本移除兼容层时踩雷。第三返回的是jnp.ndarray如果你需要转成 numpy 数组写回文件用np.asarray(samples)但要意识到这触发一次设备到主机的同步性能敏感路径上要谨慎。3.2 常用随机分布速查表JAX 的随机分布函数命名和 numpy 高度一致迁移成本比想象中小。下面是我最常用的一张表函数分布关键参数典型用途random.normal正态loc, scale参数初始化、高斯噪声random.uniform均匀minval, maxval随机裁剪、噪声random.bernoulli伯努利pDropout mask、掩码采样random.randint离散均匀minval, maxval索引采样random.categorical类别分布logits从 logits 采样、强化学习random.gumbelGumbelloc, scaleGumbel-Softmax、噪声注入random.choice/random.permutation抽样/洗牌无数据集打乱random.exponential指数无泊松过程、重采样random.poisson泊松lam计数数据random.beta/random.gamma/random.dirichlet共轭族形状参数贝叶斯采样random.multivariate_normal多维正态mean, cov潜空间采样特别提醒几乎所有分布函数都有shape参数且通常必须显式给出不写就默认生成标量。random.categorical的输入是 logits 而不是概率这是从 numpy 迁移时最容易搞混的点——直接传概率矩阵会得到错误分布必须传 log 概率。3.3 训练循环的标准模式每步一 key深度学习训练里最经典的需求是每一步都要新的 dropout mask、新的数据扰动同时整体还能复现。标准模式是每步 splitdef train_step(step_key, params, batch): key, drop_key random.split(step_key) mask random.bernoulli(drop_key, p0.3, shapeparams.shape) # ... 前向、损失、梯度 ... return key, new_params key random.key(2024) for step in range(TOTAL_STEPS): key, step_key random.split(key) key, params train_step(step_key, params, batch)这个模式的美妙之处在于只要初始 key 固定每一步的step_key都是确定可推导的。即使训练框架内部做了数据多进程切分、数据增强、dropout随机流的拓扑完全由你的 split 顺序决定不受操作系统调度或线程影响。如果你要给某个业务实体比如样本 ID绑定确定性随机数用random.fold_in(key, entity_id)而不是 split——这样以后任何时候想重算某个样本的增强版本都能精确重建。3.4 vmap 与 pmap 下的 key 分配向量化是 JAX 的看家本领随机数也要能配合。给 batch 里每个样本生成不同噪声的写法有两种先 split 再 vmap 是最稳妥的from jax import vmap keys random.split(key, batch_size) noise vmap(lambda k: random.normal(k, shape(dim,)))(keys)如果要在 pmap多设备里并行先按设备数 split再在每个设备内继续按 batch 分from jax import pmap device_keys random.split(key, jax.device_count()) def per_device(dev_key, x): sample_keys random.split(dev_key, x.shape[0]) return vmap(lambda k, xi: random.normal(k, xi.shape))(sample_keys, x) out pmap(per_device)(device_keys, x)日常使用中最常犯的错误没有按设备 split每张卡拿到同一个 key——于是所有设备生成完全相同的噪声batch 内样本多样性直接被砍掉一半训练曲线还看不出来直到你检查激活值分布才恍然大悟。规律很简单任何想要独立的维度都要先 split 出对应的 key 数量。3.5 影响性能的几个细节JAX 的随机数生成发生在设备端并且会被 XLA 编译进计算图。这意味着生成百万级样本不需要先到 CPU 绕一圈也没有 numpy 那种生成→拷贝到 GPU的传输开销。在 GPU 上random.normal(key, shape(1_000_000,))这种调用对我来说已经是随手就出结果的量级。另一个性能相关点是random.split的 num 参数在 jitted 函数里必须是静态值不能是动态 tracer。如果你写split(key, n)且 n 是函数的动态参数编译时会报错。注意random.split的 num 参数决定输出 key 的数量在 jit 下必须是静态值。解决办法是加static_argnums标记或者干脆在外部先 split 好再传进 jitted 函数。4. 确定性不止是设个种子从复现到生产环境4.1 JAX 的确定性承诺边界JAX 官方对可复现性的承诺需要精确理解在相同版本、相同硬件、相同编译配置下相同的 key 序列会产生完全一致的结果。但有三条边界必须知道。第一跨版本不保证逐位一致。JAX 团队明确表示未来版本可能修改底层算法选择或 XLA 的编译细节同样是random.key(42)生成的序列可能变化。实验复现时要锁死 JAX、XLA 和 CUDA 的版本。第二跨硬件不保证一致。GPU 上某些浮点运算的归约顺序与 CPU 不同即使随机数本身一样后续浮点累加也可能产生微小差异。第三不同 PRNG 实现Threefry 与 Philox之间天然不一致。这三条边界不是缺陷而是工程上诚实的选择——JAX 保证的是在你锁定的环境里可复现而不是全宇宙逐位一致。4.2 断点续训必须保存 key这是我在生产环境里踩过最贵的一跤。训练到第 10 万步时任务中断我重新加载了模型权重和优化器状态却忘了随机数这回事——代码里重新random.key(2024)从头派生。结果恢复训练后验证集指标莫名上涨我还高兴了一阵后来才发现只是数据顺序和 dropout mask 全变了模型根本没有真正变好。正确的做法是把 key 当作训练状态的一部分存进 checkpointckpt { params: params, opt_state: opt_state, rng: key, # 当前训练 key step: step, } # 恢复时直接继续 split key ckpt[rng] for step in range(ckpt[step], TOTAL_STEPS): key, step_key random.split(key) ...保存 key 的本质是保存随机流拓扑的当前位置。只要 key 在后续所有 dropout、增强、采样都能无缝衔接和没中断一样。这一点对于长期运行的训练任务、强化学习环境采样、在线学习系统尤为重要。4.3 与 numpy.random 的混合策略现实世界不是纯 JAX 的特别是数据加载层很多生态库还依赖np.random。我的折中策略是把 numpy 的种子也纳入 JAX 的 key 体系来管理。先用 JAX 派生一个确定性的整数种子再喂给 numpyseed_array random.randint(random.fold_in(key, dataset_id), (1,), 0, 2**31 - 1) np.random.seed(int(seed_array[0]))这样整个数据流水线的随机源头仍然是那一把 JAX key。只要固定初始 key无论 JAX 侧还是 numpy 侧都能复现同时避免两套随机体系各自为政、互相污染的混乱。如果你的代码必须同时维护多条随机流多数据集、多 worker给每条流一个独立的 fold_in 标识比在 numpy 里拷贝状态可靠得多。4.4 性能与成本为什么值得为确定性买单有人觉得显式传 key、到处 split 是麻烦不如np.random.seed()一行来得痛快。但算一笔账numpy 的全局状态让并行训练每一步都要做状态锁同步而 JAX 的 key 传递是纯数据流动编译器和运行时可以自由地并行、重排、矢量化。我在 GPU 上大规模生成随机样本的速度大概是 CPU 上 numpy 的好几倍更关键的是它不阻塞主进程。加上随机流可寻址、可重建、可断点续传这些确定性带来的工程收益替换成本通常在几周内就回本了。为确定性买单买的不是信仰是调试时间、复现能力和并行度。5. 常见问题与避坑清单5.1 高频报错速查报错信息原因解决rand() missing 1 required positional argument: key忘记传 key检查函数签名key 是第一参数PRNGKey is deprecated还在用旧 API统一改用random.key()num 类型相关报错在 jit 里传了动态 num外部先 split 或加static_argnumsfold_in 第二参数类型错误传的不是标量整数转成jnp.asarray(标量)sample shape 与预期不符忘了传 shape 或 shape 顺序写反核对每个分布签名的 shape 参数错误信息里出现Tracer在 jit/vmap 内做了不允许的动态操作把随机数生成移到计算外或调整静态参数高频错误本身不值钱值钱的是背后的习惯先检查 key 是否 split、shape 是否显式、接口是否新旧混用。我见过太多人报错后第一反应是换个 seed 试试这基本无用——错误和 seed 无关和调用姿势有关。5.2 三个让我记忆深刻的静默坑第一个固定 key 生成 dropout mask。我把 key 定义在训练循环外但忘了每步 split结果整个训练过程每一步的 dropout mask 一模一样dropout 几乎失效模型过拟合参数翻了几番。这种 bug 不报错只能靠对比两个相邻 step 的 mask 是否相同或者检查激活值分布发现。第二个把 key 也 vmap 了。我在写 per-sample 噪声时用了vmap(lambda k: random.normal(k, shape))(keys)但 keys 没 split所有样本共享同一个 key噪声完全相同。排查了很久最后发现是random.split(key, batch_size)里 batch_size 传错只 split 出 1 把。现在我的习惯是先打印 keys 的前两行确认互不相同再做 vmap。第三个跨设备复现实验时时好时坏。这是因为 GPU 上浮点归约顺序不确定而我把随机数和浮点累加混在了同一次编译里。解决办法是把随机数单独生成并固化下来比如转成np.asarray存盘浮点部分再走自己的流程两端解耦后问题消失。5.3 把 split 变成肌肉记忆的两个小工具为了根治忘记 split我自己封装了一个极简的 key 管理器class KeyManager: def __init__(self, seed): self.key random.key(seed) def next(self, n1): self.key, *subkeys random.split(self.key, n 1) return subkeys[0] if n 1 else subkeys每次要随机数就先subkey km.next()内部自动把父 key 前进。这样从流程上消灭了复用 key的可能。另一个习惯是任何函数如果内部要用随机数一律把 key 塞进签名里而不是在函数体里重新生成。这保证了函数保持纯函数既能被 jit 缓存也能被 vmapped。我在 review 别人的 JAX 代码时第一眼就找函数签名里有没有 key。5.4 生态联动从 whisper-jax 看这套设计的价值这两年 whisper-jax 把 Whisper 推理大规模批量化搬到了 JAX 上很多人惊叹于它快得离谱。但很少有人意识到JAX 生态能支撑这类高并发服务的底层原因之一正是无全局状态的设计——随机数只是其中一个缩影。批处理时每个请求、每个样本的随机操作都通过显式 key 隔离不存在互相污染分布式部署时每张卡的随机流天然独立不需要繁琐的种子同步协议。可以说你对 JAX random 的理解程度直接决定了你写的分布式训练代码是能跑还是敢上生产。最后说点个人体会。我从 numpy 迁移到 JAX 的前两周几乎每天都在和为什么两次结果一样搏斗后来有一天突然想通了——它不是要折磨你而是在用一种更诚实的姿态对待随机性随机性被看成数据流的一环而不是藏在角落的副作用。现在我再写任何训练代码都会先画一遍key 的流向图初始 key 从哪里来每个分支在哪里 splitcheckpoint 存哪把 key。这套习惯让我的实验复现率提升了不止一个档次。如果你刚开始接触 JAX别嫌随机 API 别扭把它当成一次思维升级的机会理解了 key 和 split你就理解了函数式深度学习的半边天。再分享一个小技巧每次实验把初始 seed、JAX 版本、GPU 型号记在同一个配置文件里哪怕多花十秒钟三个月后你会感激自己。