jax.nn.initializers.glorot_normal

jax.nn.initializers.glorot_normal#

jax.nn.initializers.glorot_normal(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)[source]#

构建 Glorot 正态初始化器(也称为 Xavier 正态初始化器)。

一个 Glorot 正态初始化器jax.nn.initializers.variance_scaling() 的特例,其中 scale = 1.0mode="fan_avg" 以及 distribution="truncated_normal"

参数:
  • in_axis (int | Sequence[int]) – 权重数组中输入维度对应的轴或轴序列。

  • out_axis (int | Sequence[int]) – 权重数组中输出维度对应的轴或轴序列。

  • batch_axis (Sequence[int]) – 权重数组中应被忽略的轴或轴序列。

  • dtype (DTypeLikeInexact) – 权重的 dtype。

返回:

一个初始化器。

返回类型:

Initializer

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.glorot_normal()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.41770416,  0.75262755,  0.7619329 ],
       [-0.5516644 , -0.6028657 ,  0.08661086]], dtype=float32)