jax.random.randint#

jax.random.randint(key, shape, minval, maxval, dtype=None, *, out_sharding=None)[源代码]#

以给定形状/数据类型采样 [minval, maxval) 范围内的均匀随机值。

参数:

  • key (ArrayLike) – 用作随机密钥的 PRNG 密钥。

  • shape (Shape) – 一个表示形状的非负整数元组。

  • minval (IntegerArray) – 一个整数或整数数组，可以与 shape 进行广播兼容，表示范围的最小值（包含）。

  • maxval (IntegerArray) – 一个整数或整数数组，可以与 shape 进行广播兼容，表示范围的最大值（不包含）。

  • dtype (DTypeLikeInt | None) – 可选，返回值的整数类型（如果 jax_enable_x64 为 true，则默认为 int64，否则为 int32）。

返回:

具有指定形状和 dtype 的随机数组。

返回类型:

Array

注意

randint() 使用基于模数的计算，已知在某些情况下会产生轻微的偏差。偏差的大小与 (maxval - minval) * ((2 ** nbits ) % (maxval - minval)) / 2 ** nbits 成比例：换句话说，当 (maxval - minval) 是 2 的幂时，偏差为零，否则当 (maxval - minval) 相对于采样类型的范围很小时，偏差会很小。

为了减少这种偏差，8 位和 16 位值将始终在 32 位进行采样，然后转换为请求的类型。如果您发现自己正在采样可能存在此偏差问题的数值，一种可能的替代方法是通过 uniform 进行采样。

def randint_via_uniform(key, shape, minval, maxval, dtype):
  u = jax.random.uniform(key, shape, minval=minval - 0.5, maxval=maxval - 0.5)
  return u.round().astype(dtype)

但请记住，此方法由于浮点舍入误差而存在其自身的偏差，特别是，在范围 [minval, maxval) 内可能存在一些整数无法通过此方法生成。