Equinox, 모델은 곧 Pytree

Equinox는 모델 자체가 pytree라는 다른 철학을 택해. nn.Module 같은 거 없이, 그냥 dataclass가 곧 모델이야.

pip install equinox

import equinox as eqx
import jax
import jax.numpy as jnp

# ============ 단일 Linear ============
class Linear(eqx.Module):
    weight: jnp.ndarray
    bias: jnp.ndarray

    def __init__(self, in_dim, out_dim, key):
        wkey, bkey = jax.random.split(key, 2)
        self.weight = jax.random.normal(wkey, (out_dim, in_dim)) * 0.01
        self.bias = jnp.zeros(out_dim)

    def __call__(self, x):
        return self.weight @ x + self.bias

key = jax.random.PRNGKey(0)
layer = Linear(10, 5, key)

x = jnp.zeros(10)
y = layer(x)
print(y.shape)   # (5,)

eqx.Module = @dataclass + 자동 pytree 등록. 그게 전부야.

# model 자체가 pytree
print(jax.tree.leaves(layer))
# [array(weight), array(bias)]

# tree.map 자유롭게
zeroed = jax.tree.map(jnp.zeros_like, layer)

다층 MLP

class MLP(eqx.Module):
    layers: list   # List[Linear]

    def __init__(self, dims, key):
        keys = jax.random.split(key, len(dims) - 1)
        self.layers = [Linear(d_in, d_out, k)
                       for d_in, d_out, k in zip(dims[:-1], dims[1:], keys)]

    def __call__(self, x):
        for layer in self.layers[:-1]:
            x = jax.nn.relu(layer(x))
        return self.layers[-1](x)

key = jax.random.PRNGKey(0)
model = MLP([784, 128, 64, 10], key)
y = model(jnp.zeros(784))

학습

Equinox 모델은 pytree라, 모든 JAX 변환이 그대로 작동:

def loss_fn(model, x, y):
    pred = jax.vmap(model)(x)   # batch
    return jnp.mean((pred - y) ** 2)

@jax.jit
def train_step(model, x, y, lr):
    loss, grads = jax.value_and_grad(loss_fn)(model, x, y)
    new_model = jax.tree.map(lambda p, g: p - lr * g, model, grads)
    return new_model, loss

# loop
for step in range(100):
    model, loss = train_step(model, batch_x, batch_y, 0.01)

특이점이 없고, nnx.split 같은 번거로운 절차 없이, 모델이 그냥 pytree라 jit/grad가 직접 처리해.

filter / partition, trainable과 frozen 분리

model = MLP([784, 10], key)

# 모든 param 이 변경 가능 (default)
# 일부만 학습 — eqx.filter 사용
def loss_with_frozen(diff_model, static_model, x, y):
    model = eqx.combine(diff_model, static_model)
    return loss_fn(model, x, y)

# layer 0 은 freeze, layer 1 만 학습
diff_model, static_model = eqx.partition(model,
    lambda m: True if isinstance(m, Linear) and m is model.layers[1] else False
)

grads = jax.grad(loss_with_frozen)(diff_model, static_model, x, y)

기본 제공 계층

model = eqx.nn.Sequential([
    eqx.nn.Linear(784, 128, key=k1),
    eqx.nn.Lambda(jax.nn.relu),
    eqx.nn.Linear(128, 10, key=k2),
])

# attention block
attn = eqx.nn.MultiheadAttention(
    num_heads=8, query_size=64, key=k3,
)

NNX와 Equinox 중 무엇을 고를지는 팀, 프로젝트, 코드 스타일에 달려 있어. 둘 다 프로덕션에 적용할 수 있어. JAX 코어가 같으니, 한쪽을 익히면 다른 쪽도 빠르게 따라잡을 수 있어.