Poisson Sampling

Why this matters

The Poisson distribution models count data: events occurring in a fixed time window, photon arrivals in an imaging sensor, word occurrences in a document, or queue arrivals in a server. The single parameter λ is both the mean and variance of the distribution.

In generative modelling, Poisson noise is used to simulate shot noise in scientific imaging and count-based language models. Poisson sampling is also at the heart of subsampling amplification arguments in differentially private SGD (Poisson subsampling selects each data point independently with probability q = batch_size / dataset_size).

jax.random.poisson(key, lam, shape) returns integer counts (int32 by default). Cast to float for downstream arithmetic or test contracts.

Worked mini-example

import jax, jax.numpy as jnp

key = jax.random.PRNGKey(0)

# 5 draws from Poisson(1) — low rate, mostly 0s and 1s
samples = jax.random.poisson(key, 1.0, shape=(5,))
# → int32 array of shape (5,); values typically {0, 1, 2}

# Cast to float if needed
samples_f = samples.astype(jnp.float32)

# High-rate: concentrated around lam
high = jax.random.poisson(key, 100.0, shape=(4,))
# → int32 near 100

Common pitfalls

• Returns int32, not float: cast with .astype(jnp.float32) when your downstream code or test contract expects floats.
• λ must be positive: zero gives all-zero samples; negative λ is undefined.
• shape is a tuple: shape=(int(n),) — not just int(n).
• Not a continuous distribution: confusing Poisson with Exponential (inter-arrival times) is a common conceptual error.

Problem

Implement poisson_sample(seed, lam, n) that draws n samples from Poisson(lam) and returns them as floats.

seed, lam, and n are Python scalars (floats). Return a 1-D float32 array of shape (n,) containing integer counts cast to float.

One illustrative example (not from the test set):

• poisson_sample(0, 1.0, 5.0) returns a float32 array of shape (5,) with non-negative integer values, deterministic for seed 0.