jax.random.key#

jax.random.key(seed, *, impl=None)[源代码]#

给定一个整数种子，创建一个伪随机数生成器 (PRNG) 密钥。

结果是一个包含密钥的标量数组，其 dtype 指示默认的 PRNG 实现。该实现由可选的 impl 参数确定，或者在调用此函数时由 jax_default_prng_impl 配置标志确定。

参数:

  • seed (int | ArrayLike) – 用作密钥值的 64 位或 32 位整数。

  • impl (PRNGSpecDesc | None) – 可选字符串，指定 PRNG 实现（例如 'threefry2x32'）

返回:

一个标量 PRNG 密钥数组，可供随机函数以及 split 和 fold_in 使用。

返回类型:

Array