jax.random.PRNGKey

jax.random.PRNGKey(seed)

Create a pseudo-random number generator (PRNG) key given an integer seed.

Parameters

seed (int) – a 64- or 32-bit integer used as the value of the key.

Return type

ndarray

Returns

A PRNG key, which is modeled as an array of shape (2,) and dtype uint32. The key is constructed from a 64-bit seed by effectively bit-casting to a pair of uint32 values (or from a 32-bit seed by first padding out with zeros).