Skip to content

← JAX Stochasticity step 1 of 25

Easy Primitives

Uniform Sampling

Jump to the task

Sample within an interval

A uniform distribution assigns equal probability to intervals of equal width within its range. It does not produce evenly spaced values or guarantee that a small sample covers the range well.

jax.random.uniform(key, shape, minval=low, maxval=high) requests floating-point values in the documented interval [low, high): the lower bound is included and the upper bound is excluded. If you omit the bounds, the default interval is [0, 1). This is a floating-point sampler, not a way to draw uniformly distributed integers.

One key, one vector draw

import jax

key = jax.random.PRNGKey(42)
values = jax.random.uniform(key, (3,), minval=-1.0, maxval=1.0)
# approximately [-0.022581, 0.359594, 0.232543]

The tuple (3,) requests one vector containing three samples. It does not mean three calls with the same key. Repeating the same sampler with the same key, shape, bounds and configuration repeats the vector; the key does not advance itself. Split a key when an application needs additional draws, but do not add a split to this exercise's specified construction because it changes the expected values.

The example uses float32 and the runtime's default Threefry generator. Sampling configuration is part of reproducing these numerical results.

What the bounds change

Conceptually, a value u in [0, 1) becomes low + (high - low) * u. For low = -2 and high = 6, u = 0.25 maps to 0. This explains why both the interval width and lower bound matter. Pass the bounds to uniform directly for this exercise rather than reimplementing its floating-point arithmetic.

Do not sort, normalize or evenly space the output. Return the draw in its original order. Its sample mean need not equal the midpoint of the requested interval.

Problem

Implement uniform_in_range(seed, n, low, high) by creating jax.random.PRNGKey(int(seed)) and calling jax.random.uniform once with shape (int(n),), minval=low and maxval=high.

Seed is a whole Python number in [0, 2**32), and n is a nonnegative whole Python number. Both arrive as floats; convert them to int. Low and high are finite real scalars with low < high. Return a float32 vector of shape (int(n),) using the default runtime sampling configuration. Zero samples returns an empty vector.

To JIT this interface, mark seed and n static: jax.jit(uniform_in_range, static_argnames=("seed", "n")). Low and high can remain runtime values. Calling int() on a traced seed or count does not make it static.