初始化器

flax.nnx.initializers.constant(value, dtype=<class 'jax.numpy.float64'>)

构建一个初始化器，它返回充满常数 value 的数组。

参数

• value – 用来填充初始化器的常数值。

• dtype – 可选；初始化器的默认 dtype。

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.constant(-7)
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)
Array([[-7., -7., -7.],
       [-7., -7., -7.]], dtype=float32)

flax.nnx.initializers.delta_orthogonal(scale=1.0, column_axis=-1, dtype=<class 'jax.numpy.float64'>)

构建一个用于 delta 正交核的初始化器。

参数

• scale – 均匀分布的上限。

• column_axis – 包含应该正交的列的轴。

• dtype – 权重的默认 dtype。

返回值

一个 delta 正交初始化器。传递给初始化器的形状必须是 3D、4D 或 5D。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.delta_orthogonal()
>>> initializer(jax.random.key(42), (3, 3, 3), jnp.float32)  
Array([[[ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.27858758, -0.7949833 , -0.53887904],
        [ 0.9120717 ,  0.04322892,  0.40774566],
        [-0.30085585, -0.6050892 ,  0.73712474]],

       [[ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]]], dtype=float32)

flax.nnx.initializers.glorot_normal(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 Glorot 正态初始化器（也称为 Xavier 正态初始化器）。

一个 Glorot 正态初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 1.0、mode="fan_avg" 和 distribution="truncated_normal"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> 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)

flax.nnx.initializers.glorot_uniform(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 Glorot 均匀初始化器（也称为 Xavier 均匀初始化器）。

一个 Glorot 均匀初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 1.0、mode="fan_avg" 和 distribution="uniform"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.glorot_uniform()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.50350785,  0.8088631 ,  0.81566876],
       [-0.6393332 , -0.6865721 ,  0.11003882]], dtype=float32)

flax.nnx.initializers.he_normal(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 He 正态初始化器（也称为 Kaiming 正态初始化器）。

一个 He 正态初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 2.0、mode="fan_in" 和 distribution="truncated_normal"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.he_normal()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.6604483 ,  1.1900088 ,  1.2047218 ],
       [-0.87225807, -0.95321447,  0.1369438 ]], dtype=float32)

flax.nnx.initializers.he_uniform(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 He 均匀初始化器（也称为 Kaiming 均匀初始化器）。

一个 He 均匀初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 2.0、mode="fan_in" 和 distribution="uniform"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.he_uniform()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.79611576,  1.2789248 ,  1.2896855 ],
       [-1.0108745 , -1.0855657 ,  0.17398663]], dtype=float32)

flax.nnx.initializers.kaiming_normal(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 He 正态初始化器（也称为 Kaiming 正态初始化器）。

一个 He 正态初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 2.0、mode="fan_in" 和 distribution="truncated_normal"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.he_normal()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.6604483 ,  1.1900088 ,  1.2047218 ],
       [-0.87225807, -0.95321447,  0.1369438 ]], dtype=float32)

flax.nnx.initializers.kaiming_uniform(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 He 均匀初始化器（也称为 Kaiming 均匀初始化器）。

一个 He 均匀初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 2.0、mode="fan_in" 和 distribution="uniform"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.he_uniform()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.79611576,  1.2789248 ,  1.2896855 ],
       [-1.0108745 , -1.0855657 ,  0.17398663]], dtype=float32)

flax.nnx.initializers.lecun_normal(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 Lecun 正态初始化器。

一个 Lecun 正态初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 1.0、mode="fan_in" 和 distribution="truncated_normal"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.lecun_normal()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.46700746,  0.8414632 ,  0.8518669 ],
       [-0.61677957, -0.67402434,  0.09683388]], dtype=float32)

flax.nnx.initializers.lecun_uniform(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 Lecun 均匀初始化器。

一个 Lecun 均匀初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 1.0、mode="fan_in" 和 distribution="uniform"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.lecun_uniform()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.56293887,  0.90433645,  0.9119454 ],
       [-0.71479625, -0.7676109 ,  0.12302713]], dtype=float32)

flax.nnx.initializers.normal(stddev=0.01, dtype=<class 'jax.numpy.float64'>)

构建一个初始化器，它返回实数正态分布的随机数组。

参数

• stddev – 可选；分布的标准差。

• dtype – 可选；初始化器的默认 dtype。

返回值

一个初始化器，它返回其值为均值为 0 且标准差为 stddev 的正态分布的数组。

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.normal(5.0)
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 3.0613258 ,  5.6129413 ,  5.6866574 ],
       [-4.063663  , -4.4520254 ,  0.63115686]], dtype=float32)

flax.nnx.initializers.truncated_normal(stddev=0.01, dtype=<class 'jax.numpy.float64'>, lower=-2.0, upper=2.0)

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

参数

• stddev – 可选；未截断分布的标准差。注意，此函数不会像方差缩放初始化器那样应用标准差校正，并且如果用户希望使用它，则应通过 stddev 参数自行应用此校正。

• dtype – 可选；初始化器的默认 dtype。

• lower – 表示截断下限的浮点数。在输出乘以标准差之前应用。

• upper – 表示截断上限的浮点数。在输出乘以标准差之前应用。

返回值

一个初始化器，它返回其值遵循截断正态分布的数组，其均值为 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)

flax.nnx.initializers.ones(key, shape, dtype=<class 'jax.numpy.float64'>)

一个初始化器，它返回一个全为 1 的常数数组。

参数 key 被忽略。

>>> import jax, jax.numpy as jnp
>>> jax.nn.initializers.ones(jax.random.key(42), (3, 2), jnp.float32)
Array([[1., 1.],
       [1., 1.],
       [1., 1.]], dtype=float32)

flax.nnx.initializers.ones_init()

构建一个初始化器，它返回一个全为 1 的常数数组。

>>> import jax, jax.numpy as jnp
>>> from flax.nnx import initializers
>>> ones_initializer = initializers.ones_init()
>>> ones_initializer(jax.random.key(42), (3, 2), jnp.float32)
Array([[1., 1.],
       [1., 1.],
       [1., 1.]], dtype=float32)

flax.nnx.initializers.orthogonal(scale=1.0, column_axis=-1, dtype=<class 'jax.numpy.float64'>)

构建一个初始化器，它返回均匀分布的正交矩阵。

如果形状不是方阵，则矩阵将具有正交的行或列，具体取决于哪一边更小。

参数

• scale – 均匀分布的上限。

• column_axis – 包含应该正交的列的轴。

• dtype – 权重的默认 dtype。

返回值

一个正交初始化器。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.orthogonal()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 3.9026976e-01,  7.2495741e-01, -5.6756169e-01],
       [ 8.8047469e-01, -4.7409311e-01, -1.3157725e-04]],            dtype=float32)

flax.nnx.initializers.uniform(scale=0.01, dtype=<class 'jax.numpy.float64'>)

构建一个初始化器，它返回实数均匀分布的随机数组。

参数

• scale – 可选参数；随机分布的上限。

• dtype – 可选；初始化器的默认 dtype。

返回值

一个初始化器，它返回其值在范围 [0, scale) 内均匀分布的数组。

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.uniform(10.0)
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[7.298188 , 8.691938 , 8.7230015],
       [2.0818567, 1.8662417, 5.5022564]], dtype=float32)

flax.nnx.initializers.variance_scaling(scale, mode, distribution, in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

初始化器，根据权重张量的形状调整其比例。

当 distribution="truncated_normal" 或 distribution="normal" 时，从均值为零，标准差为 \(\sqrt{\frac{scale}{n}}\) 的（截断）正态分布中抽取样本，其中 n 是

• 如果 mode="fan_in"，则为权重张量中的输入单元数量；

• 如果 mode="fan_out"，则为输出单元数量；或者

• 如果 mode="fan_avg"，则为输入和输出单元数量的平均值。

这个初始化器可以使用 in_axis、out_axis 和 batch_axis 配置，以与通用卷积或稠密层一起使用；不在这些参数中的轴被认为是“感受野”（卷积核空间轴）。

当 distribution="truncated_normal" 时，样本的绝对值在缩放之前会在 2 个标准差处截断。

当 distribution="uniform" 时，从

• 如果 dtype 为实数，则为均匀间隔，或者

• 如果 dtype 为复数，则为均匀圆盘，

中抽取样本，其均值为零，标准差为 \(\sqrt{\frac{scale}{n}}\)，其中 n 如上定义。

参数

• scale – 缩放因子（正浮点数）。

• mode – "fan_in"、"fan_out" 和 "fan_avg" 之一。

• distribution – 要使用的随机分布。 "truncated_normal"、"normal" 和 "uniform" 之一。

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

flax.nnx.initializers.xavier_normal(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 Glorot 正态初始化器（也称为 Xavier 正态初始化器）。

一个 Glorot 正态初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 1.0、mode="fan_avg" 和 distribution="truncated_normal"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> 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)

flax.nnx.initializers.xavier_uniform(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)

构建一个 Glorot 均匀初始化器（也称为 Xavier 均匀初始化器）。

一个 Glorot 均匀初始化器 是 jax.nn.initializers.variance_scaling() 的一个特例，其中 scale = 1.0、mode="fan_avg" 和 distribution="uniform"。

参数

• in_axis – 权重数组中输入维度的轴或轴序列。

• out_axis – 权重数组中输出维度的轴或轴序列。

• batch_axis – 权重数组中应该被忽略的轴或轴序列。

• dtype – 权重的 dtype。

返回值

一个初始化器。

示例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.glorot_uniform()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.50350785,  0.8088631 ,  0.81566876],
       [-0.6393332 , -0.6865721 ,  0.11003882]], dtype=float32)

flax.nnx.initializers.zeros(key, shape, dtype=<class 'jax.numpy.float64'>)

一个初始化器，它返回一个全为 0 的常数数组。

参数 key 被忽略。

>>> import jax, jax.numpy as jnp
>>> jax.nn.initializers.zeros(jax.random.key(42), (2, 3), jnp.float32)
Array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32)

flax.nnx.initializers.zeros_init()

构建一个初始化器，它返回一个全为 0 的常数数组。

>>> import jax, jax.numpy as jnp
>>> from flax.nnx import initializers
>>> zeros_initializer = initializers.zeros_init()
>>> zeros_initializer(jax.random.key(42), (2, 3), jnp.float32)
Array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32)