flax.linen.initializers.kaiming_normal#
- flax.linen.initializers.kaiming_normal(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)#
Builds a He normal initializer (aka Kaiming normal initializer).
A He normal initializer is a specialization of
jax.nn.initializers.variance_scaling()
wherescale = 2.0
,mode="fan_in"
, anddistribution="truncated_normal"
.- Parameters:
in_axis (
Union
[int
,Sequence
[int
]]) – axis or sequence of axes of the input dimension in the weights array.out_axis (
Union
[int
,Sequence
[int
]]) – axis or sequence of axes of the output dimension in the weights array.batch_axis (
Sequence
[int
]) – axis or sequence of axes in the weight array that should be ignored.dtype (
Any
) – the dtype of the weights.
- Return type:
Initializer
- Returns:
An initializer.
Example:
>>> import jax, jax.numpy as jnp >>> initializer = jax.nn.initializers.he_normal() >>> initializer(jax.random.PRNGKey(42), (2, 3), jnp.float32) Array([[ 0.6604483 , 1.1900088 , 1.2047218 ], [-0.87225807, -0.95321447, 0.1369438 ]], dtype=float32)