Getting Started
Reference Documentation
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
jit
vmap
pmap
jax.random.
rademacher
Sample from a Rademacher distribution.
key (Union[Any, PRNGKeyArray]) β a PRNG key.
Union
Any
PRNGKeyArray
shape (Sequence[int]) β The shape of the returned samples.
Sequence
int
dtype (Any) β The type used for samples.
ndarray
A jnp.array of samples, of shape shape. Each element in the output has a 50% change of being 1 or -1.