真实模型不只有一块权重，而是许多层、许多矩阵和偏置。PyTree 让 JAX 能把嵌套容器当成一个整体处理，同时仍能访问其中每个数组叶子。

1. 什么是 PyTree

PyTree 不是某个具体类，而是一类树状结构：字典、列表、元组等容器是内部节点，数组等值是叶子。

import jax
import jax.numpy as jnp

params = {
    "embedding": jnp.zeros((100, 32)),
    "block": {
        "attn": {
            "wq": jnp.zeros((32, 32)),
            "wk": jnp.zeros((32, 32)),
            "wv": jnp.zeros((32, 32)),
        },
        "mlp": [
            jnp.zeros((32, 128)),
            jnp.zeros((128, 32)),
        ],
    },
}

leaves = jax.tree.leaves(params)
print(len(leaves))

字典的键和容器结构不是模型参数；最底部的数组才是需要计算、求导和更新的叶子。

2. 为什么模型参数天然是一棵树

Transformer 的结构本来就是分层的：模型包含若干 block，每个 block 包含 attention 和 MLP，attention 又包含 Q/K/V/O 投影。PyTree 只是把这种结构直接反映到代码里。

下面是 Transformer 参数层次的结构示意，其中的 ... 表示尚未展开的参数子树，并非可直接运行的代码：

params = {
    "token_embedding": ...,          # [V, D]
    "blocks": [
        {"attn": {...}, "mlp": {...}, "norm": {...}},
        {"attn": {...}, "mlp": {...}, "norm": {...}},
    ],
    "final_norm": ...,
    "lm_head": ...,                 # [D, V]
}

3. 对整棵树做同一种操作

scaled = jax.tree.map(lambda x: x * 0.1, params)

def sgd_step(params, grads, lr):
    return jax.tree.map(lambda p, g: p - lr * g, params, grads)

两棵树一起 map 时，它们的结构必须一致。jax.grad 对参数树求导后，返回的梯度通常与参数具有相同树结构，这是优化器能统一更新所有参数的原因。

4. 统计参数量

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

估算模型规模，得到树里有多少个参数，整个参数量，矩阵的一个数值就算一个参数。比如 token embedding 的参数量是 V * D；一个不含 bias 的注意力层若有 Q/K/V/O 四个 [D, D] 矩阵，就是约 4D²；常见两层 MLP 若扩展到 4D，约为 8D²。

5. 初始化一个线性层

def init_linear(key, in_dim, out_dim, scale=0.02):
    w = scale * jax.random.normal(key, (in_dim, out_dim))
    b = jnp.zeros((out_dim,))
    return {"w": w, "b": b}

def linear(params, x):
    return x @ params["w"] + params["b"]

初始化函数返回参数树，前向函数接收参数树。这种分离让参数依赖变得显式。

6. 初始化一个 MLP

def init_mlp(key, dim, hidden_dim):
    k1, k2 = jax.random.split(key)
    return {
        "up": init_linear(k1, dim, hidden_dim),
        "down": init_linear(k2, hidden_dim, dim),
    }

def mlp(params, x):
    h = jax.nn.gelu(linear(params["up"], x))
    return linear(params["down"], h)

注意初始化时拆分了随机 key。为什么不能重复使用同一个 key，是下一篇的重点。

7. PyTree 与结构静态性

JIT 编译通常要求树结构在一次已编译程序中保持稳定。叶子的数组值可以变化，但如果今天参数字典有 w，下一步突然多出 adapter，就可能触发重新编译或结构不匹配。

训练状态也适合放进 PyTree：

state = {
    "params": params,
    "step": jnp.array(0),
    "optimizer": {"momentum": ...},
}

其中超参数、函数、字符串等不一定适合当动态数组叶子。成熟框架会提供 dataclass 注册和静态字段机制。

8. 检查树的实用函数

shape_tree = jax.tree.map(lambda x: (x.shape, x.dtype), params)
finite_tree = jax.tree.map(lambda x: jnp.all(jnp.isfinite(x)), params)

调试大模型时，不要打印所有权重；打印 shape、dtype、均值、标准差、最大绝对值和是否包含 NaN，信息更有效。

小练习

对参数树的每个叶子计算 L2 范数。

norm_tree = jax.tree.map(lambda x: jnp.linalg.norm(x), params)

跟之前 shape_tree/finite_tree 的写法完全一致，只是 lambda 换成了计算 L2 范数的 jnp.linalg.norm。