jax.random.bernoulli#

jax.random.bernoulli(key,p=0.5,shape=None,mode='low',*,out_sharding=None)[source]#

Sample Bernoulli random values with given shape and mean.

The values are distributed according to the probability mass function:

\[f(k; p) = p^k(1 - p)^{1 - k}\]

where\(k \in \{0, 1\}\) and\(0 \le p \le 1\).

Parameters:

• key (ArrayLike) – a PRNG key used as the random key.

• p (RealArray) – optional, a float or array of floats for the mean of the randomvariables. Must be broadcast-compatible withshape. Default 0.5.

• shape (Shape |None) – optional, a tuple of nonnegative integers representing the resultshape. Must be broadcast-compatible withp.shape. The default (None)produces a result shape equal top.shape.

• mode (str) – optional, “high” or “low” for how many bits to use when sampling.default=’low’. Set to “high” for correct sampling at small values ofp. When sampling in float32, bernoulli samples with mode=’low’ produceincorrect results for p < ~1E-7. mode=”high” approximately doubles thecost of sampling.

Returns:

A random array with boolean dtype and shape given byshape ifshapeis not None, or elsep.shape.

Return type:

Array