jax.nn.initializers.truncated_normal#

jax.nn.initializers.truncated_normal(stddev=0.01, dtype=<class 'numpy.float64'>, lower=-2.0, upper=2.0)[source]#

构建一个返回截断正态随机数组的初始化器。

参数:
  • stddev (RealNumeric) – 可选;未截断分布的标准差。 请注意,此函数不应用方差缩放初始化器中完成的标准差校正,如果用户希望使用它,则应通过 stddev 参数自行应用此校正。

  • dtype (DTypeLikeInexact) – 可选;初始化器的默认 dtype。

  • lower (RealNumeric) – 浮点数,表示截断的下限。 在输出乘以 stddev 之前应用。

  • upper (RealNumeric) – 浮点数,表示截断的上限。 在输出乘以 stddev 之前应用。

返回:

一个初始化器,返回值的数组遵循均值为 0 和标准差为 stddev 的截断正态分布,范围为 \(\rm{lower * stddev} < x < \rm{upper * stddev}\)

返回类型:

初始化器

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.truncated_normal(5.0)
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 2.9047365,  5.2338114,  5.29852  ],
       [-3.836303 , -4.192359 ,  0.6022964]], dtype=float32)