flâneur

jax.random module — JAX documentation

jax.readthedocs.io · 1,694 words · saved by 1 readers

The jax.random package provides a number of routines for deterministic generation of sequences of pseudorandom numbers. Unlike the stateful pseudorandom number generators (PRNGs) that users of NumPy and SciPy may be accustomed to, JAX random functions all require an explicit PRNG state to be passed as a first argument. The random state is described by a special array element type that we call a key, usually generated by the jax.random.key() function: This key can then be used in any of JAX’s random number generation routines: Note that using a key does not modify it, so reusing the same key will lead to the same result: If you need a new random number, you can use jax.random.split() to generate new subkeys: Note Typed key arrays, with element types such as key above, were introduced in JAX v0.4.16. Before then, keys were conventionally represented in uint32 arrays, whose final dimension represented the key’s bit-level representation. Both forms of key array can still be created an

jax.random module # Utilities for pseudo-random number generation. The jax.random package provides a number of routines for deterministic generation of sequences of pseudorandom numbers. Basic usage # >>> seed = 1701 >>> num_steps = 100 >>> key = jax . random . key ( seed ) >>> for i in range ( num_steps ): ... key , subkey = jax . random . split ( key ) ... params = compiled_update ( subkey , params , next ( batches )) PRNG keys # Unlike the stateful pseudorandom number generators (PRNGs) that users of NumPy and SciPy may be accustomed to, JAX random functions all require an explicit PRNG sta

related reading