Skip to content

← JAX Stochasticity step 24 of 25

Medium Primitives

Poisson Sampling

Jump to the task

Draw a count for a fixed interval

A Poisson distribution models nonnegative integer counts. Its parameter lam is the expected count in the interval you are modeling. For example, if arrivals follow a constant-rate Poisson process at two per minute, a three-minute interval has lam = 6. A particular interval can still contain fewer or more than six arrivals.

The distribution has mean lam and variance lam, so its standard deviation is sqrt(lam). These are population properties, not exact statistics of a short vector of samples. Not every kind of count data follows a Poisson model: that assumption needs to fit the application.

A fractional mean still produces whole counts

Lam can be fractional. A mean of 0.5 does not produce half an event; it produces counts such as 0, 1 or 2 with different probabilities. For any positive lam, the distribution allows arbitrarily large nonnegative counts, although large counts may be very unlikely.

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(42)
counts = jax.random.poisson(key, 5.0, shape=(4,))
# [7, 4, 3, 2] with the exercise runtime's default settings
result = counts.astype(jnp.float32)
# [7., 4., 3., 2.]

This requests four counts in one call. The float conversion meets the exercise's output format; integer arrays already support arithmetic. It does not turn the distribution into a continuous one. In a Poisson process, the exponential distribution describes waiting times, while the Poisson distribution describes counts in a fixed interval.

The term Poisson subsampling in privacy literature refers to independently including each record with a Bernoulli decision. It is not an instruction to call this Poisson count sampler to select the records.

Reproduce the requested stream

Use one seed-derived key for the vector draw. Repeating the same key, rate, shape and sampling configuration repeats the output. Splitting the key first or substituting a rounded normal draw changes the task. The examples use the runtime's default Threefry generator and Poisson sampling settings.

Problem

Implement poisson_sample(seed, lam, n) by creating jax.random.PRNGKey(int(seed)), drawing with jax.random.poisson and shape=(int(n),), then converting the counts to float32.

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. Lam is a finite, nonnegative floating-point scalar. Return shape (int(n),). Lam zero gives all-zero counts; zero samples gives an empty vector.

For JIT, mark seed and n static: jax.jit(poisson_sample, static_argnames=("seed", "n")). Lam can remain a runtime value. int() cannot make a traced seed or output length static.