jax.extend.random.threefry_prng_impl

jax.extend.random.threefry_prng_impl#

jax.extend.random.threefry_prng_impl = ((2,), <function threefry_seed>, <function threefry_split>, <function threefry_random_bits>, <function threefry_fold_in>, 'threefry2x32', 'fry')#

指定PRNG键的形状和操作。

PRNG 实现由一个键类型 K 和一组对这些键进行操作的函数确定。键类型 K 是一个元素类型为 uint32 的数组类型,其形状由 key_shape 指定。每个操作的类型签名是:

seed :: int[] -> K
fold_in :: K -> int[] -> K
split[shape] :: K -> K[*shape]
random_bits[shape, bit_width] :: K -> uint<bit_width>[*shape]

PRNG 实现通过 PRNGKeyArray 类适应于键 K 的类似数组的对象,该类应通过 random_seed 函数创建。