jax.nn.initializers.glorot_uniform#

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

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

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

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.glorot_uniform()
>>> initializer(jax.random.PRNGKey(42), (2, 3), jnp.float32)  
Array([[ 0.50350785,  0.8088631 ,  0.81566876],
       [-0.6393332 , -0.6865721 ,  0.11003882]], dtype=float32)