불변성: 배열을 직접 바꾸지 않고 a = a.at[0].set(99)처럼 새 배열을 만든다.
float32 기본 dtype: x64를 활성화하지 않으면 부동소수점 기본값은 float32다.
키 기반 난수: 별도 모듈인 jax.random에서 PRNGKey와 명시적인 키를 사용한다.
자동 장치 배치: 배열과 연산이 사용 가능한 가속기에 자동으로 배치된다.
⚠️ 자주 만나는 버그
NumPy의 전역 seed 방식 코드를 그대로 JAX로 옮기면 같은 키를 반복해서 사용해 매번 같은 난수가 나올 수 있고, 그 결과 학습이 진행되지 않을 수 있어. jax.random.PRNGKey와 split을 중심으로 난수 흐름을 다시 설계해야 해.
기능 | NumPy | JAX
-------------------|-------------------|------------------
mutate | a[0] = 1 | a = a.at[0].set(1)
default float | float64 | float32
random | np.random (전역) | jax.random + key
device | RAM 만 | CPU/GPU/TPU 자동
Code
import numpy as np
import jax.numpy as jnp
# NumPy: in-place mutation works fine
np_arr = np.array([1, 2, 3])
np_arr[0] = 99
print(np_arr) # [99, 2, 3]
# JAX: in-place mutation raises an error
jnp_arr = jnp.array([1, 2, 3])
# jnp_arr[0] = 99 # ERROR: JAX arrays are immutable
# Instead, use .at[].set() to create a NEW array
new_arr = jnp_arr.at[0].set(99)
print(new_arr) # [99, 2, 3]
print(jnp_arr) # [1, 2, 3] — original unchanged!
import jax.numpy as jnp
x = jnp.array([10, 20, 30, 40, 50])
# Set a value
x_new = x.at[2].set(99) # [10, 20, 99, 40, 50]
# Add to a value
x_add = x.at[2].add(5) # [10, 20, 35, 40, 50]
# Multiply
x_mul = x.at[2].mul(2) # [10, 20, 60, 40, 50]
# Slice updates
x_slice = x.at[1:3].set(0) # [10, 0, 0, 40, 50]
# Conditional update with jnp.where
mask = x > 25
x_where = jnp.where(mask, x * 2, x) # [10, 20, 60, 80, 100]
import jax
import jax.numpy as jnp
# Check available devices
print(jax.devices()) # e.g., [CudaDevice(id=0)] or [TpuDevice(id=0)]
print(jax.default_backend()) # 'gpu', 'tpu', or 'cpu'
# Arrays are created on the default device
x = jnp.array([1.0, 2.0, 3.0])
print(x.devices()) # Shows which device(s) the array is on
# Explicitly place on a device
cpu_device = jax.devices('cpu')[0]
x_cpu = jax.device_put(x, cpu_device)
import numpy as np
import jax.numpy as jnp
x = jnp.array([1.0, 2.0, 3.0])
# This works but BYPASSES JIT — the array is silently moved to CPU
result = np.sum(x) # Uses NumPy, not JAX — no JIT, no GPU
# Always use jnp for JAX arrays
result = jnp.sum(x) # Correct — uses JAX, can be JIT-compiled
a[mask] = 0처럼 in-place 변경을 사용하는 다섯 줄짜리 NumPy 코드를 준비해. JAX에서 (1) jnp.where, (2) .at[mask].set(0), (3) 별도의 함수형 helper라는 세 방식으로 다시 작성해. 셋 모두 jit 안에서 문제없이 컴파일되는지 확인하고 코드를 저장해. 이후에도 반복해서 쓰게 될 패턴이야.
Progress
Progress is local-only — sign in to sync across devices.