jax.random.fold_in函数介绍

jax.random.fold_in 是一个用于将一个值与随机数生成器的状态合并的函数。它通常用于在分布式计算或并行计算中，将多个随机数生成器的状态合并为一个。

import jax
import jax.numpy as jnp

# 需要混合的数据
key = jax.random.PRNGKey(42)
data = 1

# 调用 jax.random.fold_in 函数
new_key = jax.random.fold_in(key, data)
# 打印结果
print(new_key)


# 多个数据
data_batch = jnp.array([1, 2, 3])

# 随机数生成器状态的 fold_in 函数
def fold_in_state(data):
    return jax.random.fold_in(key, data)

# 使用 jax.vmap 对一批数据进行随机数生成器状态的混合
new_keys = jax.vmap(fold_in_state)(data_batch)

# 打印结果
print(new_keys)

参考：

https://jax.readthedocs.io/en/latest/_autosummary/jax.random.fold_in.html#