추적 깊이 보기: 추상 값과 ShapedArray

추적이 어떻게 일어나는지 들여다보면 JAX의 모든 동작이 자연스러워져.

jit 함수를 처음 호출하면 인자가 실제 배열이 아니라 Tracer 객체로 대체돼. Tracer는 배열과 비슷하지만, 구체적인 값 대신 추상적인 메타데이터만 가져:

shape, 모양 (예: (32, 768))
dtype, 자료형(예: float32)
그게 전부야.

이 메타데이터를 ShapedArray 라고 불러. 함수 안에서 일어나는 모든 연산은 ShapedArray 위에서, 실제 계산이 아니라, "어떤 op이 어떤 shape/dtype을 만들지" 추적해.

import jax

@jax.jit
def f(x):
    print(type(x))   # 첫 호출에서 한 번만 출력. 무엇이 찍힐까?
    print(x)
    return jnp.sum(x ** 2)

f(jnp.arange(5.))
# <class 'jax.interpreters.partial_eval.JaxprTracer'>
# Traced<ShapedArray(float32[5])>with<DynamicJaxprTrace(level=1/0)>

Python의 print는 추적 시점에 한 번만 도니까, 여기서 보이는 게 Tracer 객체의 정체. x는 5개짜리 float32배열이라는 건 알지만, 값은 몰라.

이게 왜 중요하냐면 Tracer의 한계 때문에 Python의 일부 동작이 안 돼:

@jax.jit
def f(x):
    if x > 0:        # ❌ Tracer 의 truth value 가 없음
        return x
    return -x

# 해결
@jax.jit
def f(x):
    return jnp.where(x > 0, x, -x)  # value 사용 안 함, abstract op

Tracer가 가진 정보로 가능한 것:

x.shape, OK (정적)
x.dtype, OK
x.ndim, OK
jnp.sum(x), OK (op 추가, abstract)
x[0], OK if index is 정적

안 되는 것:

x.item(), 구체적인 값 필요
int(x), float(x), 구체적인 값
if x > 0:, 실제 불리언 값 필요
x.numpy(), concrete

대부분의 함수는 추적할 수 있어, 추상 값으로도 op 흐름은 같으니까. 그래서 거의 모든 NumPy-style 코드가 그대로 jit이 가능해.

🔬 Jaxpr 보기

jax.make_jaxpr(f)(x) 호출하면 JAX가 추적한 IR을 인쇄. 컴파일 안 하고 추적만. 함수의 op 흐름이 정확히 보여. 디버깅 + 성능 분석에 유용. print(jax.make_jaxpr(f)(x)), 한 번 해 봐.

import jax

def f(x):
    return jnp.sum(x ** 2 + jnp.sin(x))

print(jax.make_jaxpr(f)(jnp.arange(5.)))
# { lambda ; a:f32[5]. let
#     b:f32[5] = integer_pow[y=2] a
#     c:f32[5] = sin a
#     d:f32[5] = add b c
#     e:f32[] = reduce_sum[axes=(0,)] d
#   in (e,) }

읽는 법: let <var>:<type> = <op>[<params>] <args>. 5개짜리 float32가 들어와서 power, sin, add가 차례로 적용되고 reduce_sum으로 합쳐져 0-d float32 (scalar)가 나와. 컴파일 단계에서는 이 IR이 XLA에 넘어가.

이 사고 모형 한 번 잡으면 JAX의 거의 모든 오류 메시지를 해석할 수 있어.

Code

import jax
import jax.numpy as jnp

@jax.jit
def show_tracer(x):
    # During tracing, x is a Tracer wrapping a ShapedArray
    print(f"Type during trace: {type(x)}")
    # Shape is known:
    print(f"Shape: {x.shape}")
    # But values are NOT known:
    # if x[0] > 0: ...  # This would fail!
    return x + 1

show_tracer(jnp.array([1.0, 2.0, 3.0]))
# Prints something like:
# Type during trace: <class 'jax...JaxprTracer'>
# Shape: (3,)

import jax
import jax.numpy as jnp

def my_fn(x, y):
    z = jnp.sin(x) + jnp.cos(y)
    return jnp.sum(z)

# See the traced computation
jaxpr = jax.make_jaxpr(my_fn)(jnp.ones(3), jnp.ones(3))
print(jaxpr)
# { lambda ; a:f32[3] b:f32[3]. let
#     c:f32[3] = sin a
#     d:f32[3] = cos b
#     e:f32[3] = add c d
#     f:f32[] = reduce_sum[axes=(0,)] e
#   in (f,) }

🔬 Jaxpr 보기

Code

External links

Exercise

Progress

댓글 0