ott.solvers.nn.models.ICNN.init_fn#

ICNN.init_fn(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.PRNGKey(42), (2, 3), jnp.float32)  
Array([[ 3.0613258 ,  5.6129413 ,  5.6866574 ],
       [-4.063663  , -4.4520254 ,  0.63115686]], dtype=float32)