jax.nn.initializers.normal

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

Builds an initializer that returns real normally-distributed random arrays.

Parameters:

• stddev (Any) – optional; the standard deviation of the distribution.

• dtype (Any) – optional; the initializer’s default dtype.

Return type:

Initializer

Returns:

An initializer that returns arrays whose values are normally distributed with mean 0 and standard deviation 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)