jax.random.bernoulli¶

jax.random.bernoulli(key, p=0.5, shape=None)[source]¶

Sample Bernoulli random values with given shape and mean.

Parameters

key (ndarray) – a PRNGKey used as the random key.
p (ndarray) – optional, a float or array of floats for the mean of the random variables. Must be broadcast-compatible with shape. Default 0.5.
shape (Optional[Sequence[int]]) – optional, a tuple of nonnegative integers representing the result shape. Must be broadcast-compatible with p.shape. The default (None) produces a result shape equal to p.shape.

Return type

ndarray

Returns

A random array with boolean dtype and shape given by shape if shape is not None, or else p.shape.