jax.nn.initializers.glorot_uniform
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()
wherescale = 1.0
,mode="fan_avg"
, anddistribution="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)