如何在 JAX 的 vmap 中正确处理 PRNG 分割？-Python教程-PHP中文网

如何在 JAX 的 vmap 中正确处理 PRNG 分割？

本文详解如何在 JAX 中为带随机噪声的时序模拟（如 SDE 求解）安全地管理伪随机数生成器（PRNG）状态，重点解决 fori_loop 内部去相关采样、vmap 批量并行时的键分割与广播问题。

本文详解如何在 jax 中为带随机噪声的时序模拟（如 sde 求解）安全地管理伪随机数生成器（prng）状态，重点解决 `fori_loop` 内部去相关采样、`vmap` 批量并行时的键分割与广播问题。

在 JAX 中，PRNG 状态不可变且必须显式传递——这是函数式编程范式的核心约束，也是避免隐式随机性耦合、确保可复现性的关键。当你在 simulate 中使用 fori_loop 迭代求解随机微分方程（SDE）时，若对每一步都复用同一 key，jax.random.multivariate_normal 将产生完全相同的噪声序列，导致所有时间步退化为确定性演化；而若在循环内盲目调用 jax.random.split(key) 并仅取一个子键，则未使用的子键将丢失，破坏键空间的可预测性与可复现性。

✅ 正确做法是：在循环中使用 jax.random.fold_in(key, i) 为每个时间步生成唯一、确定性、无冲突的子密钥。fold_in 将整数索引 i 混入原始 key，生成语义上独立的新密钥，既避免了显式分割-丢弃的资源浪费，又保证了跨步噪声的统计独立性（只要原始 key 足够随机）。修改后的核心函数如下：

import jax
import jax.numpy as jnp

# 假设 dt, d, covariance_matrix, sqrt_dt 已定义
def evolve(u, t, key):
    noise = jax.random.multivariate_normal(
        key, 
        jnp.zeros(d), 
        covariance_matrix, 
        shape=(1,)  # 注意：shape=(1,) 对应单步噪声向量
    ).squeeze(0)  # 去除 batch 维度
    return u + dt * b(t, u) + sigma(t, u) * sqrt_dt * noise

def simulate(x, t, key):
    k = jnp.floor(t / dt).astype(int)
    # 使用 fold_in 为每步生成独立 key
    def body_fn(i, u):
        step_key = jax.random.fold_in(key, i)
        return evolve(u, i * dt, step_key)
    u = jax.lax.fori_loop(0, k, body_fn, x)
    return u

⚠️ 注意事项：

fold_in 是纯函数，不消耗原始 key，因此无需返回更新后的 key；
evolve 不应自行 split 输入 key（除非内部需多路采样），外部控制密钥派生逻辑更清晰、更易调试；
若 evolve 内部需多个独立噪声（如不同维度独立扰动），可在该函数内对 step_key 调用一次 split 并分配子键，但务必全部使用，避免“密钥泄漏”。

当扩展至批量模拟时，vmap 要求输入张量在映射轴上形状一致。此时，不应将单个 key 直接传入 vmap（JAX 会尝试广播它，但所有批次将共享同一密钥，导致完全相关的噪声样本）。正确方式是：预先将初始 key 分割为 batch_size 个独立子密钥，并与 x_batch 和 t_batch 同轴对齐：

def batch_simulate(x_batch, t_batch, key):
    assert x_batch.shape[0] == t_batch.shape[0]
    # 将单个 key 分割为 batch_size 个独立密钥
    key_batch = jax.random.split(key, x_batch.shape[0])
    # vmap 沿第 0 轴并行调用 simulate
    return jax.vmap(simulate)(x_batch, t_batch, key_batch)

# 使用示例
key = jax.random.PRNGKey(42)
x_batch = jnp.array([[1.0, 0.0], [0.0, 1.0]])  # shape: (2, d)
t_batch = jnp.array([1.0, 2.0])                 # shape: (2,)
results = batch_simulate(x_batch, t_batch, key)

? 总结关键原则：