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)