jax.random.wrap_key_data

jax.random.wrap_key_data(key_bits_array, *, impl=None)

将密钥数据位数组包装到 PRNG 密钥数组中。

参数：

• key_bits_array (Array) – uint32 数组，其尾部形状对应于由 impl 指定的 PRNG 实现的密钥形状。

• impl (PRNGSpecDesc | None | None) – 可选，指定 PRNG 实现，如 random.key。

返回：

一个 PRNG 密钥数组，其数据类型是 jax.dtypes.prng_key 的子类型

对应于 impl，其形状等于 key_bits_array.shape 在密钥位维度之前的领先形状。