Getting Started
Reference Documentation
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
Create a pseudo-random number generator (PRNG) key given an integer seed.
The resulting key carries the default PRNG implementation, as determined by the jax_default_prng_impl config flag.
jax_default_prng_impl
seed (int) – a 64- or 32-bit integer used as the value of the key.
int
Union[Any, PRNGKeyArray]
Union
Any
PRNGKeyArray
A PRNG key, consumable by random functions as well as split and fold_in.
split
fold_in