순수성을 깨뜨리는 일곱 가지 패턴

순수성이 깨지는 지점을 미리 알면 디버깅이 빨라져. 다음 일곱 가지 함정을 차례로 살펴보자.

1. global 또는 nonlocal 값 읽기

scale = 2.0

@jax.jit
def f(x):
    return x * scale  # ❌ scale 이 첫 호출 때 capture 됨, 이후 변경 무시

scale을 9.0으로 바꾸고 다시 호출해도 함수는 계속 2.0을 곱해. 첫 추적에서 기록한 값이 캐시에 남기 때문이야.

2. list나 dict 같은 변경 가능한 컨테이너 갱신하기

history = []

@jax.jit
def f(x):
    history.append(x)  # ❌ side effect. trace 시점에 한 번 append.
    return x * 2

3. print, 로깅, 파일 쓰기

@jax.jit
def f(x):
    print(f"x is {x}")  # ❌ trace 때만 print 됨. 게다가 x 는 Tracer object!
    return x ** 2

일반 print는 추적할 때 한 번 실행되므로 추적 과정을 살펴볼 때는 쓸모가 있어. 하지만 함수가 실행될 때마다 출력된다고 생각하면 안 돼. 실제 실행 시점의 값을 출력하려면 jax.debug.print('{x}', x=x)를 사용해.

4. random.random 또는 np.random 사용하기

import random

@jax.jit
def f(x):
    noise = random.gauss(0, 1)  # ❌ trace 때 한 번. 매 호출마다 같은 noise.
    return x + noise

JAX에서는 jax.random과 명시적인 키로 난수를 전달해야 해. Track 8에서 자세히 배울 거야.

5. 예외로 제어 흐름 만들기

@jax.jit
def f(x):
    try:
        return jnp.log(x)
    except:  # ❌ trace 시점엔 실제 값 없음. exception 안 일어남.
        return jnp.zeros_like(x)

6. time.time()이나 datetime.now() 읽기

@jax.jit
def f(x):
    seed = int(time.time())  # ❌ trace 한 순간의 시간만 capture
    ...

7. iterator 또는 generator 상태 사용하기

it = iter(range(100))

@jax.jit
def f(x):
    return x + next(it)  # ❌ iterator state 가 mutate. 그것도 첫 호출 한 번.

⚠️ 오류 없이 틀릴 수 있어

이 패턴들은 JAX가 반드시 오류를 내는 문제가 아니야. 이전에 추적해 캐시한 결과를 조용히 다시 사용할 수 있어. 학습이 진행되지 않거나 결과가 이상하다면 순수성부터 의심해. 시험 삼아 일반 print를 함수 안에 넣고 같은 형태의 입력으로 두 번 호출했는데 두 번째에는 출력되지 않는다면, 추적과 실제 실행을 구분해 다시 살펴봐야 해.

방어 방법은 외부 상태를 모두 명시적인 입력과 출력으로 바꾸는 거야.

# global 대신 인자로
def f(x, scale):  # ✅
    return x * scale

# state 는 in/out 으로
def step(state, x):
    new_state = state + 1
    return new_state, x * new_state

# random 은 key 로
def f(x, key):
    noise = jax.random.normal(key, x.shape)
    return x + noise

# print 는 jax.debug.print
@jax.jit
def f(x):
    jax.debug.print("x is {x}", x=x)  # ✅ runtime print
    return x ** 2

이 일곱 가지 함정을 기억하면 "왜 실행되지 않지?"라는 문제의 90%를 훨씬 빠르게 좁힐 수 있어.

Code

import jax.numpy as jnp

x = jnp.array([1, 2, 3])
# x[0] = 99  # TypeError: JAX arrays are immutable

# Fix: use .at[].set()
x_new = x.at[0].set(99)  # Returns new array, x is unchanged

import jax
import jax.numpy as jnp

learning_rate = 0.01

# BAD: reads from closure
@jax.jit
def update_bad(params, grads):
    return params - learning_rate * grads

# GOOD: pass as argument
@jax.jit
def update_good(params, grads, lr):
    return params - lr * grads

# OR: use static_argnums for values that rarely change
@jax.jit
def update_static(params, grads, lr):
    return params - lr * grads
# JAX will recompile when lr changes, but that's acceptable if it rarely does

import jax
import jax.numpy as jnp

@jax.jit
def fn_with_print(x):
    print("This runs during TRACING only, not execution!")
    y = x + 1
    print(f"y = {y}")  # Prints a tracer object, not a number
    return y

result = fn_with_print(jnp.array(5.0))
# Output during first call:
# "This runs during TRACING only, not execution!"
# "y = Traced<ShapedArray(float32[])>with<DynamicJaxprTrace...>"

# Second call: no print at all — JIT reuses the cached trace
result2 = fn_with_print(jnp.array(10.0))

@jax.jit
def fn_with_debug_print(x):
    y = x + 1
    jax.debug.print("y = {}", y)  # Prints at execution time!
    return y

fn_with_debug_print(jnp.array(5.0))   # Prints "y = 6.0"
fn_with_debug_print(jnp.array(10.0))  # Prints "y = 11.0"

import jax
import jax.numpy as jnp

# NumPy uses global state — IMPURE
import numpy as np
np.random.seed(42)
a = np.random.randn(3)  # Mutates global RNG state
b = np.random.randn(3)  # Different result — depends on hidden state

# JAX uses explicit keys — PURE
key = jax.random.PRNGKey(42)
key1, key2 = jax.random.split(key)
a = jax.random.normal(key1, (3,))  # Deterministic given key1
b = jax.random.normal(key2, (3,))  # Deterministic given key2

# Same key always gives same result
a_again = jax.random.normal(key1, (3,))
print(jnp.allclose(a, a_again))  # True — pure!

import jax
import jax.numpy as jnp

# PROBLEMATIC under JIT: Python if depends on a traced value
@jax.jit
def bad_relu(x):
    if x > 0:  # ConcretizationTypeError!
        return x
    else:
        return 0.0

# GOOD: use JAX control flow
@jax.jit
def good_relu(x):
    return jnp.where(x > 0, x, 0.0)

⚠️ 오류 없이 틀릴 수 있어

Code

External links

Exercise

Progress

댓글 0