Getting Started
Reference Documentation
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
jit
vmap
pmap
jax.random.
split
Splits a PRNG key into num new keys by adding a leading axis.
key (Union[Any, PRNGKeyArray]) – a PRNG key (from PRNGKey, split, fold_in).
Union
Any
PRNGKeyArray
PRNGKey
fold_in
num (int) – optional, a positive integer indicating the number of keys to produce (default 2).
int
Union[Any, PRNGKeyArray]
An array-like object of num new PRNG keys.