Skip to content

← JAX Stochasticity step 21 of 25

Easy Primitives

Beta Distribution Sampling

Jump to the task

Sample a fraction between zero and one

A beta distribution describes a random fraction or probability. Its two positive shape parameters, alpha and beta, control where values tend to fall. A beta draw is a real-valued fraction; a Bernoulli draw is a yes-or-no outcome. They are different kinds of random variables.

The mean is alpha / (alpha + beta). The sum alpha + beta controls concentration when their ratio is fixed. For example, Beta(2, 2) and Beta(20, 20) both have mean 0.5, but Beta(20, 20) has less spread. The mean is not a promise about a small batch's average.

Parameters Distribution shape
alpha = beta = 1 Uniform density
Both below 1 U-shaped, with density rising near both ends
Both above 1 One interior peak
alpha = 2, beta = 5 More weight toward zero; mean 2/7
alpha = 5, beta = 2 More weight toward one; mean 5/7

When both parameters exceed one, the mode is (alpha - 1) / (alpha + beta - 2). It need not equal the mean: Beta(2, 5) has mode 1/5 and mean 2/7.

A concrete draw

import jax

key = jax.random.PRNGKey(0)
values = jax.random.beta(key, 1.0, 1.0, shape=(4,))
# approximately [0.203631, 0.895658, 0.931155, 0.774002]

Beta(1, 1) has a uniform distribution, but this does not make four samples evenly spread. It also does not make this call return the same values as jax.random.uniform with the same key. Matching a distribution is different from matching a sampler's random stream.

In Bayesian modeling, Beta(a, b) is a conjugate prior for a Bernoulli success probability: s successes and f failures give Beta(a+s, b+f). A uniform Beta(1, 1) prior is one modeling choice, not a universal statement of no prior knowledge. Beta draws also serve as mixing weights, where values near zero or one favor one of two inputs.

Numerical boundaries

The ideal continuous distribution puts probability one strictly between zero and one. Floating-point sampling can round to endpoints, especially for small shape parameters. Do not clip the sampler's output to force a strict interval. Invalid nonpositive parameters are outside this exercise's contract; do not rely on a particular error or NaN.

Problem

Implement beta_sample(seed, alpha, beta_param, n) by creating jax.random.PRNGKey(int(seed)) and calling jax.random.beta once with alpha, beta_param and shape=(int(n),).

Seed is a whole Python number in [0, 2**32), supplied as a float. N is a nonnegative whole Python number supplied as a float. Both shape parameters are finite positive floats. Return a float32 vector of shape (int(n),); zero samples returns an empty vector.

Use the runtime's default Threefry generator and beta sampling settings without splitting the key or replacing the sampler. For JIT, mark seed and n static: jax.jit(beta_sample, static_argnames=("seed", "n")). Alpha and beta_param can remain runtime values. int() cannot turn a traced seed or output length into a static value.