Pytree Utility: flatten, unflatten, 디버깅

pytree를 리프 목록으로 풀었다가 다시 묶을 수 있어. 디버깅, 직렬화, 임의 조작에 꼭 필요한 기능이야.

import jax

params = {
    "W1": jnp.zeros((3, 4)),
    "b1": jnp.zeros(4),
    "W2": jnp.zeros((4, 2)),
    "b2": jnp.zeros(2),
}

# 1. flatten — leaves + structure
leaves, treedef = jax.tree.flatten(params)
# leaves: [W1, b1, W2, b2]  (list of arrays)
# treedef: 트리 구조 정보 (PyTreeDef)

print(f"leaves: {len(leaves)}")
for i, l in enumerate(leaves):
    print(f"  [{i}] shape={l.shape}")

# 2. unflatten — leaves 와 structure 로 재구성
restored = jax.tree.unflatten(treedef, leaves)
# restored 는 원본 params 와 동일

# 3. tree.leaves / tree.structure (단축)
just_leaves = jax.tree.leaves(params)
just_struct = jax.tree.structure(params)

활용 1: 모든 매개변수를 한 vector로

def params_to_flat_vector(params):
    '''pytree → 1D vector'''
    leaves = jax.tree.leaves(params)
    return jnp.concatenate([l.ravel() for l in leaves])

def flat_vector_to_params(flat, template):
    '''1D vector → pytree (template 의 구조 사용)'''
    leaves = jax.tree.leaves(template)
    sizes = [l.size for l in leaves]
    shapes = [l.shape for l in leaves]
    treedef = jax.tree.structure(template)

    chunks = jnp.split(flat, jnp.cumsum(jnp.array(sizes))[:-1])
    new_leaves = [c.reshape(s) for c, s in zip(chunks, shapes)]
    return jax.tree.unflatten(treedef, new_leaves)

# 사용
flat = params_to_flat_vector(params)
print(flat.shape)   # (총 parameter 개수,)
restored = flat_vector_to_params(flat, params)

2nd-order 옵티마이저 (LBFGS, Newton's 메서드)가 pytree를 flat vector로 다뤄야 할 때 사용해. jax.flatten_util.ravel_pytree가 같은 일을 한 줄로 해 줌:

from jax.flatten_util import ravel_pytree

flat, unravel = ravel_pytree(params)
# flat: 1D vector
# unravel: 함수, flat → 원래 pytree

# 사용
restored = unravel(flat)
new_flat = flat + 0.01   # 어떤 transform
new_params = unravel(new_flat)

활용 2: 디버깅, tree 구조 출력

def print_tree(tree, prefix=""):
    if isinstance(tree, dict):
        for k, v in tree.items():
            if isinstance(v, dict):
                print(f"{prefix}{k}/")
                print_tree(v, prefix + "  ")
            else:
                print(f"{prefix}{k}: shape={v.shape}, dtype={v.dtype}")

print_tree(params)
# W1: shape=(3, 4), dtype=float32
# b1: shape=(4,), dtype=float32
# ...

# 또는 한 줄 — leaves 와 path 함께
import jax.tree_util
def display(p, v):
    return f"{'.'.join(str(x.key) for x in p)}: {v.shape}"

jax.tree_util.tree_map_with_path(display, params)

활용 3: 매개변수 count

def count_params(params):
    return sum(x.size for x in jax.tree.leaves(params))

n = count_params(transformer_params)
print(f"{n:,} parameters")
# 예: 124,440,000 (124M model)

# layer 별
for path, leaf in jax.tree_util.tree_leaves_with_path(params):
    print(f"{path}: {leaf.size:,}")

🔧 flatten/unflatten의 정신

pytree가 매력적인 이유, 두 가지를 분리해. 한 쪽에는 "구조" (어떤 모양으로 중첩되어 있나), 한 쪽에는 "값" (실제 배열). 같은 구조에 다른 값을 끼워 넣을 수 있어, 체크포인트 적재, 매개변수 교체 같은 모든 조작이 깔끔해져.

주의, treedef 자체는 hashable이라 jit의 static_argnames에 넣을 수 있어. 그래서 추적 안에서 트리 구조가 자동으로 캐시돼.

Code

import jax
import jax.numpy as jnp

params = {
    'layer1': {'w': jnp.array([[1.0, 2.0], [3.0, 4.0]]), 'b': jnp.array([0.1, 0.2])},
    'layer2': {'w': jnp.array([[5.0], [6.0]]), 'b': jnp.array([0.3])},
}

# Flatten: get leaves and structure separately
leaves, treedef = jax.tree.flatten(params)
print(f"Number of leaves: {len(leaves)}")  # 4
print(f"Tree structure: {treedef}")

# You can reconstruct the original tree
params_rebuilt = treedef.unflatten(leaves)
# params_rebuilt == params (same structure and values)

# Useful: concatenate all params into one vector
all_params = jnp.concatenate([x.ravel() for x in leaves])
print(f"Total params as one vector: {all_params.shape}")  # (9,)

# Get paths alongside leaves
path_leaves, treedef = jax.tree.flatten_with_path(params)

for path, leaf in path_leaves:
    path_str = jax.tree_util.keystr(path)
    print(f"{path_str}: shape={leaf.shape}, dtype={leaf.dtype}")

# Output:
# ['layer1']['b']: shape=(2,), dtype=float32
# ['layer1']['w']: shape=(2, 2), dtype=float32
# ['layer2']['b']: shape=(1,), dtype=float32
# ['layer2']['w']: shape=(2, 1), dtype=float32

# Find NaN parameters
def check_for_nans(params):
    path_leaves, _ = jax.tree.flatten_with_path(params)
    for path, leaf in path_leaves:
        if jnp.any(jnp.isnan(leaf)):
            print(f"NaN found at {jax.tree_util.keystr(path)}!")

# Check parameter statistics
def param_summary(params):
    path_leaves, _ = jax.tree.flatten_with_path(params)
    for path, leaf in path_leaves:
        name = jax.tree_util.keystr(path)
        print(f"{name}: mean={leaf.mean():.4f}, std={leaf.std():.4f}, "
              f"shape={leaf.shape}")

param_summary(params)
# ['layer1']['b']: mean=0.1500, std=0.0500, shape=(2,)
# ['layer1']['w']: mean=2.5000, std=1.1180, shape=(2, 2)
# ...

# PyTorch equivalent: model.state_dict()
# state_dict = model.state_dict()
# for name, param in state_dict.items():
#     print(f"{name}: {param.shape}")
# Output: "layer1.weight: torch.Size([4, 3])"

# JAX equivalent using pytrees:
path_leaves, _ = jax.tree.flatten_with_path(params)
for path, leaf in path_leaves:
    print(f"{jax.tree_util.keystr(path)}: {leaf.shape}")
# Output: "['layer1']['w']: (3, 4)"

# Both give named access to all parameters — different syntax, same idea

🔧 flatten/unflatten의 정신

Code

External links

Exercise

Progress

댓글 0