One of the most powerful patterns in JAX is combining randomness with vmap for vectorized random sampling. Instead of writing a loop to generate random values for a batch, you split keys and vmap over them.
import jax
import jax.numpy as jnp
def sample_and_transform(key):
"""Generate one random sample and apply some transformation."""
x = jax.random.normal(key, (3,))
return jnp.sin(x) + x ** 2
# Generate 1000 independent samples using vmap
key = jax.random.key(42)
keys = jax.random.split(key, 1000)
results = jax.vmap(sample_and_transform)(keys)
print(results.shape) # (1000, 3)This pattern is essential for Monte Carlo methods, data augmentation, and anything that requires batched randomness. Every element in the batch gets a unique key, so the samples are truly independent.
Dropout with vmap
def dropout(x, key, rate=0.5):
mask = jax.random.bernoulli(key, 1.0 - rate, x.shape)
return jnp.where(mask, x / (1.0 - rate), 0.0)
# Apply different dropout masks to each sample in a batch
def forward_single(params, x, key):
h = jax.nn.relu(x @ params['w1'] + params['b1'])
k1, k2 = jax.random.split(key)
h = dropout(h, k1, rate=0.3)
out = h @ params['w2'] + params['b2']
return out
# vmap over both inputs and keys
batch_forward = jax.vmap(forward_single, in_axes=(None, 0, 0))
# Each sample gets its own dropout mask
key = jax.random.key(0)
batch_keys = jax.random.split(key, 32) # one key per sample
# predictions = batch_forward(params, batch_x, batch_keys)The Partitionable PRNG (Since JAX 0.5.0)
Since JAX 0.5.0 (February 2025), the jax_threefry_partitionable config option is enabled by default. This means JAX's default Threefry PRNG is auto-parallelizable — random number generation works efficiently across multiple devices without any extra code from you.
# Before JAX 0.5.0: random ops could be slow under pmap/sharding
# because the PRNG wasn't designed for partitioning.
# Since JAX 0.5.0: partitionable by default!
# Random ops automatically shard across devices efficiently.
# No config changes needed.
# Important: the partitionable PRNG produces DIFFERENT values
# than the old non-partitionable version for the same seed.
# jax.random.key(42) gives different numbers in JAX 0.5+ vs 0.4.x💡 Why This Matters
The partitionable PRNG was a major quality-of-life improvement for distributed training. Before it, random operations under pmap or sharding could silently produce correlated randomness across devices, or force expensive all-gather operations. Now it "just works" — each device computes its own independent random numbers in parallel. If you're upgrading from pre-0.5.0, be aware that your random values will change for the same seeds.
Practical Example: Data Augmentation
def augment_image(key, image):
"""Apply random augmentations to a single image."""
k1, k2, k3 = jax.random.split(key, 3)
# Random horizontal flip
flip = jax.random.bernoulli(k1)
image = jnp.where(flip, jnp.flip(image, axis=1), image)
# Random brightness adjustment
brightness = jax.random.uniform(k2, minval=0.8, maxval=1.2)
image = image * brightness
# Random crop offset (for a simple center crop with jitter)
offset = jax.random.randint(k3, (2,), 0, 8)
image = jax.lax.dynamic_slice(image, (*offset, 0),
(224, 224, 3))
return jnp.clip(image, 0.0, 1.0)
# Augment entire batch with independent randomness
batch_augment = jax.vmap(augment_image)
key = jax.random.key(0)
batch_keys = jax.random.split(key, 64) # 64 images
# augmented_batch = batch_augment(batch_keys, image_batch)