Model.set_seed()

Model.set_seed(seed)

Splits and sets the seed / PRNG key.

Parameters:

seed (PRNGKeyArray) – The seed is split and distributed to the seed nodes of the model. Must be a KeyArray, i.e. an array of shape (2,) and dtype uint32. See jax.random for more details.

Return type:

Model