vmap 안에서 난수가 어떻게 동작하는지, 매우 중요. 잘못 짜면 배치의 모든 예제가 같은 노이즈를 받아서 학습이 깨져.
흔한 함정
def add_noise(x, key):
return x + random.normal(key, x.shape) * 0.1
x_batch = jnp.zeros((32, 10)) # 32 examples
key = random.PRNGKey(0)
# 잘못 — 모든 example 이 같은 noise
batched = jax.vmap(add_noise, in_axes=(0, None))
y = batched(x_batch, key)
# 32 examples 다 똑같은 noise pattern. bug!
올바른 패턴, 각 예제마다 다른 키:
def add_noise(x, key):
return x + random.normal(key, x.shape) * 0.1
# 32 개의 독립적 key
keys = random.split(key, 32)
# vmap 으로 — 각 example 이 자기 key
batched = jax.vmap(add_noise, in_axes=(0, 0))
y = batched(x_batch, keys)
# 32 examples 다 다른 noise!
새 JAX의 PRNG는 배열 모양으로 직접 처리가능해. random.normal(key, (32, 10)) 한 번이 split 후 32번 호출과 통계적으로 동등하지만 더 효율.
# 이 둘은 다른 결과지만 둘 다 OK
# 방식 1: 한 key, big shape
big = random.normal(key, (32, 10)) # 32 example × 10 feature 의 noise
# 방식 2: split, vmap
keys = random.split(key, 32)
small = jax.vmap(lambda k: random.normal(k, (10,)))(keys)
방식 1이 더 빠르고 메모리 효율도 좋아. 그러나 함수가 단일 예제 단위로 작성되어 있고 vmap으로 batch 처리하는 패턴에선 방식 2가 자연스러워.
vmap의 각 배치 원소가 독립적인 난수를 사용하려면 각자에게 독립적 키를 줘야 한다. 한 키를 broadcast 하면 모두가 같은 난수를 봐 (가끔 그게 의도지만 거의 항상 버그). split → vmap, 또는 한 번에 큰 shape으로 샘플링.
실전 디버깅 팁: 학습이 이상하면 첫 번째와 마지막 예제의 배치 결과를 출력해 봐. 같으면 난수 시드 버그. 다르면 난수는 OK, 다른 곳을 봐야 해.
Code
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)
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)
# 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
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)