JAX PyTree 自定义数据结构的注册方法

如果一个 Python 自定义类同时保存模型数组和少量配置，直接把它传给 JAX 变换通常会遇到类型无法解释的问题。原因是：未注册的自定义类默认只是一个叶子，JAX 看不到其中的数组。注册 PyTree 的关键不是给类加一个标记，而是明确回答两个问题：哪些字段要递归参与计算，哪些字段只用于定义结构并保持静态。

官方地址：https://docs.jax.dev/en/latest/pytrees.html

最常用的注册思路是：把数组、标量或嵌套树放进 children，把可哈希且不会随计算变化的配置放进 aux_data；随后由 unflatten 根据两者重新构造对象。

先判断自定义类为什么没有被展开

JAX 内置认识列表、元组、字典等容器，但普通自定义类如果没有加入注册表，会整体作为一个 leaf。下面的 kind 是固定配置，features 和 labels 才是要交给 JAX 变换的数组。

import jax
import jax.numpy as jnp

class SampleBatch:
    def __init__(self, features, labels, kind):
        # features 和 labels 是动态数据；kind 是结构元数据。
        self.features = features
        self.labels = labels
        self.kind = kind

batch = SampleBatch(jnp.ones((2, 3)), jnp.zeros(2), "train")
# 未注册时，JAX 会把整个对象当成一个叶子，而不是展开三个属性。
print(jax.tree.leaves(batch))

用 register_pytree_node 把拆分规则写在类外

当类来自第三方库或不方便修改类定义时，可以调用 jax.tree_util.register_pytree_node。flatten 函数返回 (children, aux_data)：children 会继续被 JAX 递归遍历，aux_data 会写入树结构并参与缓存键，因此必须是稳定、可哈希的静态值。

import jax
import jax.numpy as jnp

class SampleBatch:
    def __init__(self, features, labels, kind):
        self.features = features
        self.labels = labels
        self.kind = kind

def flatten_batch(value):
    # 数组进入 children，kind 进入可哈希的静态辅助数据。
    children = (value.features, value.labels)
    aux_data = value.kind
    return children, aux_data

def unflatten_batch(aux_data, children):
    # children 已被 JAX 还原为新的数组或 tracer，再重新组装对象。
    features, labels = children
    return SampleBatch(features, labels, aux_data)

jax.tree_util.register_pytree_node(
    SampleBatch, flatten_batch, unflatten_batch
)

@jax.jit
def mean_loss(value):
    # 注册完成后，JIT 可以看到对象内部的数组字段。
    return jnp.mean(value.features) + jnp.mean(value.labels)

batch = SampleBatch(jnp.ones((2, 3)), jnp.zeros(2), "train")
print(mean_loss(batch))

这里的 unflatten_batch 不应只返回 children，也不能把 kind 忘掉。JAX 在树映射、梯度或 JIT 重建参数时都会调用它；只要重建出的对象仍保留同样的字段边界，外部类就能像普通 PyTree 一样工作。

类自己维护规则时使用 register_pytree_node_class

如果可以修改类定义，装饰器形式更容易把注册规则和数据结构放在一起。tree_flatten 与上面的 flatten 函数职责相同，tree_unflatten 接收静态元数据和已还原的 children。

@jax.tree_util.register_pytree_node_class
class ModelInput:
    def __init__(self, tokens, mask, mode):
        # mode 会影响分支选择，因此作为静态元数据保存。
        self.tokens = tokens
        self.mask = mask
        self.mode = mode

    def tree_flatten(self):
        # 只把会参与数组运算的字段交给 JAX 递归处理。
        children = (self.tokens, self.mask)
        aux_data = self.mode
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        # aux_data 和 children 的顺序必须与 tree_flatten 对应。
        tokens, mask = children
        return cls(tokens, mask, aux_data)

@jax.jit
def encode(value):
    # 静态 mode 可用于 Python 分支，数组字段保持动态。
    if value.mode == "masked":
        return value.tokens * value.mask
    return value.tokens

value = ModelInput(jnp.arange(3), jnp.array([1, 0, 1]), "masked")
print(encode(value))

这种方式适合领域模型、训练批次和配置对象。要特别注意：aux_data 不能放入 JAX array、NumPy array 或会变化的列表；如果某个字段本来是动态数组，却被放进静态元数据，JIT 缓存键和重编译行为都会变得不合理。

dataclass 用 data_fields 和 meta_fields 声明边界

已有 dataclass 时，可以使用 register_dataclass。最直观的兼容写法是显式列出动态字段和静态字段：

from dataclasses import dataclass
from functools import partial
import jax
import jax.numpy as jnp

@partial(
    jax.tree_util.register_dataclass,
    data_fields=["x", "y"],
    meta_fields=["mode"],
)
@dataclass
class Pair:
    x: jax.Array
    y: jax.Array
    mode: str

@jax.jit
def combine(value):
    # mode 是静态字符串，x 和 y 是动态数组。
    if value.mode == "add":
        return value.x + value.y
    return value.x - value.y

print(combine(Pair(jnp.arange(3), jnp.ones(3), "add")))

data_fields 应是数组、标量或其叶子可被 JAX 处理的嵌套树；meta_fields 必须是静态、可哈希、不可变的对象。较新的 JAX 也支持省略字段列表并用 jax.tree.static() 标注静态 dataclass 字段，但显式写出字段边界更容易在多个环境之间复用。

注册完成后用四个问题复查

1. flatten 是否只返回动态内容：数组和需要递归遍历的嵌套树放 children，固定字符串、整数或枚举值放 aux_data/meta_fields。
2. 静态值能否稳定哈希：不要把 JAX array、NumPy array 或可变容器当作静态配置。
3. unflatten 是否完整：字段数量、顺序和类型要能让对象恢复到可继续传入 jax.jit 的状态。
4. JIT 分支是否真的需要静态值：如果只是数组计算，不必把动态数据提升成静态元数据；否则可能增加缓存变体。

简单记忆即可：register_pytree_node 适合外部注册，register_pytree_node_class 适合类内实现，register_dataclass 适合按字段声明。三者解决的是同一个问题——让 JAX 知道如何拆开和重新组装你的自定义数据结构。

常见问题

为什么注册后仍然提示类型不能参与 JIT？

通常是注册发生在调用之后，或 flatten 返回的 children 中仍包含无法被 JAX 处理的对象。确保注册代码在构造实例和调用变换前执行，并检查动态字段的叶子。

配置字段应该放 children 还是 aux_data？

会随每次计算变化、需要被数组变换的值放 children；用于选择固定 Python 分支且可哈希的配置放 aux_data 或 meta_fields。