jax.nn.initializers.xavier_normal#

jax.nn.initializers.xavier_normal(in_axis=-2, out_axis=-1, batch_axis=(), dtype=None)[source]#

Builds a Glorot normal initializer (aka Xavier normal initializer).

A Glorot normal initializer is a specialization of jax.nn.initializers.variance_scaling() where scale = 1.0, mode="fan_avg", and distribution="truncated_normal".

Parameters:

  • in_axis (int | Sequence[int]) – axis or sequence of axes of the input dimension in the weights array.

  • out_axis (int | Sequence[int]) – axis or sequence of axes of the output dimension in the weights array.

  • batch_axis (int | Sequence[int]) – axis or sequence of axes in the weight array that should be ignored.

  • dtype (DTypeLikeInexact | None) – the dtype of the weights.

Returns:

An initializer.

Return type:

Initializer

Examples:

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