Model.set_seed()

Contents

Model.set_seed()#

Model.set_seed(seed)[source]#

Splits and sets the seed / PRNG key.

Parameters:: seed (Array) – The seed is split and distributed to the seed nodes of the model. Must be a jax RNG key array that satisfies jnp.issubdtype(key.dtype, jax.dtypes.prng_key). See jax.random and https://docs.jax.dev/en/latest/jep/9263-typed-keys.html for more details.
Return type:: Model