JAX의 난수 분포

JAX가 제공하는 확률분포 집합은 NumPy와 거의 같지만 모든 함수에 키 인자가 추가돼.

import jax
from jax import random

key = random.PRNGKey(0)

# 기본 분포
random.uniform(key, shape, minval=0, maxval=1)
random.normal(key, shape)            # 표준 정규분포
random.bernoulli(key, p=0.5, shape=())
random.exponential(key, shape)
random.gamma(key, a, shape)
random.poisson(key, lam, shape)
random.chisquare(key, df, shape)
random.t(key, df, shape)
random.beta(key, a, b, shape)
random.gumbel(key, shape)
random.laplace(key, shape)
random.logistic(key, shape)
random.cauchy(key, shape)
random.dirichlet(key, alpha, shape)

# 이산 분포
random.categorical(key, logits, shape=())
random.choice(key, a, shape, replace=True, p=None)

# permutation
random.permutation(key, x, axis=0, independent=False)

# 정수
random.randint(key, shape, minval, maxval)
random.bits(key, shape, dtype=jnp.uint32)

대부분, NumPy와 동일한 인자 + 첫 인자가 키야.

실전 사용 예

# 1. 가중치 초기화 — Xavier/Glorot
key = random.PRNGKey(0)
fan_in, fan_out = 100, 50
limit = jnp.sqrt(6 / (fan_in + fan_out))
W = random.uniform(key, (fan_in, fan_out), minval=-limit, maxval=limit)

# 또는 He
W_he = random.normal(key, (fan_in, fan_out)) * jnp.sqrt(2 / fan_in)

# 2. Dropout
def dropout(x, key, p=0.5):
    mask = random.bernoulli(key, 1 - p, x.shape)
    return x * mask / (1 - p)

# 3. Sampling from softmax
logits = jnp.array([1.0, 2.0, 0.5, 3.0])
sample = random.categorical(key, logits, shape=(10,))   # (10,) 의 indices

# 4. Gumbel-softmax (differentiable sampling)
g = random.gumbel(key, logits.shape)
soft_sample = jax.nn.softmax((logits + g) / temperature)

# 5. Permutation
indices = jnp.arange(100)
shuffled = random.permutation(key, indices)

NumPy 호환성

NumPy 코드의 난수 call을 옮길 때, 키 인자만 추가:

np.random.normal(0, 1, (10,))      → random.normal(key, (10,))  (loc=0, scale=1)
np.random.uniform(low=0, high=1)    → random.uniform(key, ...)
np.random.choice(a, size=10)        → random.choice(key, a, shape=(10,))
np.random.permutation(x)            → random.permutation(key, x)

주의: NumPy의 일부 분포는 JAX에선 약간 다른이름이거나 인자 순서 달라. 새로운 분포 처음 쓸 때 docs 한 번 확인해.

💡 normal은 standard만

JAX의 random.normal은 평균 0, 표준편차 1만 지원해. 다른 평균 / std라면 직접 transform: mu + std * random.normal(key, shape). NumPy처럼 normal(loc, scale) 인자를 직접 받는 방식은 아니야. 한 번 깜박하면 헷갈려.

multivariate normal도 별도 함수 (random.multivariate_normal), 인자 형태가 자세함. 큰 공분산 행렬은 Cholesky 분해 비용이 크므로 주의해.

Code

import jax
import jax.numpy as jnp

key = jax.random.key(42)
keys = jax.random.split(key, 8)

# Uniform [0, 1)
uniform = jax.random.uniform(keys[0], shape=(3,))
# [0.299, 0.784, 0.033]

# Normal (mean=0, std=1)
normal = jax.random.normal(keys[1], shape=(3,))
# [-0.272, 1.085, -0.533]

# Bernoulli (coin flips)
coins = jax.random.bernoulli(keys[2], p=0.7, shape=(5,))
# [True, True, False, True, True]

# Categorical (sample from discrete distribution)
logits = jnp.array([1.0, 2.0, 0.5])
category = jax.random.categorical(keys[3], logits, shape=(4,))
# [1, 1, 0, 1] — category 1 most likely

# Truncated normal (clipped to [lower, upper])
trunc = jax.random.truncated_normal(keys[4], lower=-2.0, upper=2.0, shape=(3,))

# Randint (random integers)
ints = jax.random.randint(keys[5], shape=(3,), minval=0, maxval=10)

# Permutation (shuffling)
shuffled = jax.random.permutation(keys[6], jnp.arange(5))

# Exponential
exp_samples = jax.random.exponential(keys[7], shape=(3,))

# NumPy                              # JAX
# np.random.randn(3, 4)             jax.random.normal(key, (3, 4))
# np.random.rand(3, 4)              jax.random.uniform(key, (3, 4))
# np.random.randint(0, 10, (3,))    jax.random.randint(key, (3,), 0, 10)
# np.random.choice(arr, size=5)     jax.random.choice(key, arr, (5,))
# np.random.shuffle(arr)            jax.random.permutation(key, arr)

# Key difference: JAX always needs a key, and never mutates input

# Xavier/Glorot initialization from scratch
def glorot_normal(key, shape):
    fan_in, fan_out = shape[-2], shape[-1]
    std = jnp.sqrt(2.0 / (fan_in + fan_out))
    return jax.random.normal(key, shape) * std

key = jax.random.key(0)
w = glorot_normal(key, (256, 128))
print(f"Mean: {w.mean():.4f}, Std: {w.std():.4f}")
# Mean: ~0.0, Std: ~0.072

💡 normal은 standard만

Code

External links

Exercise

Progress

댓글 0