jax.random.PRNGKey创建伪随机数生成器密钥

jax.random.PRNGKey 是 JAX 库中用于创建伪随机数生成器密钥（PRNG key）的函数。
PRNG key 是 JAX 中用于生成伪随机数序列的关键要素，通过分割（splitting）可以生成新的 PRNG key，确保生成的随机数序列是不相关的。

import jax
import jax.numpy as jnp
import haiku as hk

# 创建一个 PRNG key
rng_key = jax.random.PRNGKey(42)

print("rng_key")
print(rng_key)
print(type(rng_key))

# 生成 [0, 1) 范围内的随机数
random_numbers = jax.random.uniform(rng_key, shape=(3, 3))

print(random_numbers)
# 分割 PRNG key，生成新的 PRNG key
new_rng_key1, new_rng_key2 = jax.random.split(rng_key, num=2)

print(new_rng_key1)  
print(new_rng_key2)

random_numbers_1 = jax.random.uniform(new_rng_key1, shape=(3, 3))
random_numbers_2 = jax.random.uniform(new_rng_key2, shape=(3, 3))

print(random_numbers_1)
print(random_numbers_2)

new_rng_key, dropout_rng_key = hk.next_rng_key()

print(new_rng_key)
print(dropout_rng_key)