自定义 pytree 节点

本节将解释如何通过使用 jax.tree_util.register_pytree_node() 和 jax.tree.map()，在 JAX 中扩展被视为 pytree 内部节点（pytree 节点）的 Python 类型集合。

为什么需要这样做？在之前的示例中，pytree 展示为列表、元组和字典，而其他所有内容都被视为 pytree 叶子。这是因为如果你定义了自己的容器类，除非你向 JAX 注册它，否则它将被视为 pytree 叶子。即使你的容器类内部包含树结构，情况也是如此。例如：

import jax

class Special(object):
  def __init__(self, x, y):
    self.x = x
    self.y = y

jax.tree.leaves([
    Special(0, 1),
    Special(2, 4),
])

[<__main__.Special at 0x77843a4c9280>, <__main__.Special at 0x778469e15a00>]

相应地，如果你尝试使用 jax.tree.map() 并期望叶子是容器内部的元素，你将会收到一个错误：

jax.tree.map(lambda x: x + 1,
  [
    Special(0, 1),
    Special(2, 4)
  ])

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[2], line 1
----> 1 jax.tree.map(lambda x: x + 1,
      2   [
      3     Special(0, 1),
      4     Special(2, 4)

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/tree.py:156, in map(f, tree, is_leaf, *rest)
    116 def map(f: Callable[..., Any],
    117         tree: Any,
    118         *rest: Any,
    119         is_leaf: Callable[[Any], bool] | None = None) -> Any:
    120   """Maps a multi-input function over pytree args to produce a new pytree.
    121 
    122   Args:
   (...)    154     - :func:`jax.tree.reduce`
    155   """
--> 156   return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/tree_util.py:394, in tree_map(f, tree, is_leaf, *rest)
    392 leaves, treedef = tree_flatten(tree, is_leaf)
    393 all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
--> 394 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/tree_util.py:394, in <genexpr>(.0)
    392 leaves, treedef = tree_flatten(tree, is_leaf)
    393 all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
--> 394 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

Cell In[2], line 1, in <lambda>(x)
----> 1 jax.tree.map(lambda x: x + 1,
      2   [
      3     Special(0, 1),
      4     Special(2, 4)

TypeError: unsupported operand type(s) for +: 'Special' and 'int'

作为解决方案，JAX 允许通过全局类型注册表来扩展被视为内部 pytree 节点的类型集合。此外，注册类型的数值会被递归遍历。

首先，使用 jax.tree_util.register_pytree_node() 注册一个新类型：

from jax.tree_util import register_pytree_node

class RegisteredSpecial(Special):
  def __repr__(self):
    return "RegisteredSpecial(x={}, y={})".format(self.x, self.y)

def special_flatten(v):
  """Specifies a flattening recipe.

  Params:
    v: The value of the registered type to flatten.
  Returns:
    A pair of an iterable with the children to be flattened recursively,
    and some opaque auxiliary data to pass back to the unflattening recipe.
    The auxiliary data is stored in the treedef for use during unflattening.
    The auxiliary data could be used, for example, for dictionary keys.
  """
  children = (v.x, v.y)
  aux_data = None
  return (children, aux_data)

def special_unflatten(aux_data, children):
  """Specifies an unflattening recipe.

  Params:
    aux_data: The opaque data that was specified during flattening of the
      current tree definition.
    children: The unflattened children

  Returns:
    A reconstructed object of the registered type, using the specified
    children and auxiliary data.
  """
  return RegisteredSpecial(*children)

# Global registration
register_pytree_node(
    RegisteredSpecial,
    special_flatten,    # Instruct JAX what are the children nodes.
    special_unflatten   # Instruct JAX how to pack back into a `RegisteredSpecial`.
)

现在你可以遍历特殊的容器结构：

jax.tree.map(lambda x: x + 1,
  [
   RegisteredSpecial(0, 1),
   RegisteredSpecial(2, 4),
  ])

[RegisteredSpecial(x=1, y=2), RegisteredSpecial(x=3, y=5)]

或者，你可以在类中定义适当的 tree_flatten 和 tree_unflatten 方法，并使用 register_pytree_node_class() 对其进行装饰：

from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class RegisteredSpecial2(Special):
  def __repr__(self):
    return "RegisteredSpecial2(x={}, y={})".format(self.x, self.y)

  def tree_flatten(self):
    children = (self.x, self.y)
    aux_data = None
    return (children, aux_data)

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    return cls(*children)


def show_example(structured):
  flat, tree = structured.tree_flatten()
  unflattened = RegisteredSpecial2.tree_unflatten(tree, flat)
  print(f"{structured=}\n  {flat=}\n  {tree=}\n  {unflattened=}")


show_example(RegisteredSpecial2(1., 2.))

structured=RegisteredSpecial2(x=1.0, y=2.0)
  flat=(1.0, 2.0)
  tree=None
  unflattened=RegisteredSpecial2(x=1.0, y=2.0)

现代 Python 配备了许多有用的工具，使得定义容器变得更加容易。有些可以与 JAX 直接配合使用，但另一些则需要格外小心。

例如，Python 的 NamedTuple 子类不需要注册即可被视为 pytree 节点类型：

from typing import NamedTuple, Any

class MyOtherContainer(NamedTuple):
  name: str
  a: Any
  b: Any
  c: Any

# NamedTuple subclasses are handled as pytree nodes, so
# this will work out-of-the-box.
jax.tree.leaves([
    MyOtherContainer('Alice', 1, 2, 3),
    MyOtherContainer('Bob', 4, 5, 6)
])

['Alice', 1, 2, 3, 'Bob', 4, 5, 6]

请注意，name 字段现在显示为叶子，因为所有元组元素都是子项。这就是当你不需要费力注册该类时发生的情况。

在定义展平函数（unflattening functions）时，通常 children 应该包含数据结构中的所有动态元素（数组、动态标量和 pytree），而 aux_data 应该包含将被整合到 treedef 结构中的所有静态元素。JAX 有时需要比较 treedef 以判断其是否相等，或者计算其哈希值以用于 JIT 缓存，因此必须确保指定的辅助数据支持有意义的哈希和相等性比较。

与 NamedTuple 子类不同，使用 @dataclass 装饰的类不会自动成为 pytree。但是，可以使用 jax.tree_util.register_dataclass() 装饰器将它们注册为 pytree：

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

@functools.partial(jax.tree_util.register_dataclass,
                   data_fields=['a', 'b', 'c'],
                   meta_fields=['name'])
@dataclass
class MyDataclassContainer(object):
  name: str
  a: Any
  b: Any
  c: Any

# MyDataclassContainer is now a pytree node.
jax.tree.leaves([
  MyDataclassContainer('apple', 5.3, 1.2, jnp.zeros([4])),
  MyDataclassContainer('banana', np.array([3, 4]), -1., 0.)
])

[5.3, 1.2, Array([0., 0., 0., 0.], dtype=float32), array([3, 4]), -1.0, 0.0]

请注意，name 字段不会显示为叶子。这是因为我们将其包含在 jax.tree_util.register_dataclass() 的 meta_fields 参数中，表明它应被视为元数据/辅助数据，就像上面 RegisteredSpecial 中的 aux_data 一样。现在 MyDataclassContainer 的实例可以传递给 JIT 编译的函数，并且 name 将被视为静态的（有关静态参数的更多信息，请参阅 将参数标记为静态）。

@jax.jit
def f(x: MyDataclassContainer | MyOtherContainer):
  return x.a + x.b

# Works fine! `mdc.name` is static.
mdc = MyDataclassContainer('mdc', 1, 2, 3)
y = f(mdc)

将其与 MyOtherContainer（NamedTuple 子类）对比。由于 name 字段是一个 pytree 叶子，JIT 期望它能转换为 jax.Array，因此以下代码会引发错误：

moc = MyOtherContainer('moc', 1, 2, 3)
y = f(moc)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[9], line 2
      1 moc = MyOtherContainer('moc', 1, 2, 3)
----> 2 y = f(moc)

    [... skipping hidden 3 frame]

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/pjit.py:649, in _infer_input_type(fun, dbg_fn, explicit_args)
    647   dbg = dbg_fn()
    648   arg_description = f"path {dbg.arg_names[i] if dbg.arg_names is not None else 'unknown'}"
--> 649   raise TypeError(
    650     f"Error interpreting argument to {fun} as an abstract array."
    651     f" The problematic value is of type {type(x)} and was passed to"
    652     f" the function at {arg_description}.\n"
    653     "This typically means that a jit-wrapped function was called with a non-array"
    654     " argument, and this argument was not marked as static using the"
    655     " static_argnums or static_argnames parameters of jax.jit."
    656   ) from None
    657 if config.mutable_array_checks.value:
    658   check_no_aliased_ref_args(dbg_fn, avals, explicit_args)

TypeError: Error interpreting argument to <function f at 0x7784322271a0> as an abstract array. The problematic value is of type <class 'str'> and was passed to the function at path x.name.
This typically means that a jit-wrapped function was called with a non-array argument, and this argument was not marked as static using the static_argnums or static_argnames parameters of jax.jit.

用于操作 pytree 的整套函数都在 jax.tree_util 中。

自定义 pytree 及使用意外值的初始化

用户自定义 pytree 对象的另一个常见陷阱是，JAX 转换有时会使用意外的值来初始化它们，导致在初始化时进行的任何输入验证都可能失败。例如：

class MyTree:
  def __init__(self, a):
    self.a = jnp.asarray(a)

register_pytree_node(MyTree, lambda tree: ((tree.a,), None),
    lambda _, args: MyTree(*args))

tree = MyTree(jnp.arange(5.0))

jax.vmap(lambda x: x)(tree)      # Error because object() is passed to `MyTree`.

<__main__.MyTree at 0x778439e59a60>

jax.jacobian(lambda x: x)(tree)  # Error because MyTree(...) is passed to `MyTree`.

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[11], line 1
----> 1 jax.jacobian(lambda x: x)(tree)  # Error because MyTree(...) is passed to `MyTree`.

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/api.py:830, in jacrev.<locals>.jacfun(*args, **kwargs)
    823 f = lu.wrap_init(
    824     fun, kwargs,
    825     debug_info=debug_info(
    826         "jacrev", fun, args, kwargs,
    827         static_argnums=(argnums,) if isinstance(argnums, int) else argnums))
    828 f_partial, dyn_args = argnums_partial(f, argnums, args,
    829                                       require_static_args_hashable=False)
--> 830 tree_map(partial(_check_input_dtype_jacrev, holomorphic, allow_int), dyn_args)
    831 if has_aux:
    832   y, pullback, aux = _vjp(f_partial, *dyn_args, has_aux=True)

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/tree_util.py:394, in tree_map(f, tree, is_leaf, *rest)
    392 leaves, treedef = tree_flatten(tree, is_leaf)
    393 all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
--> 394 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

Cell In[10], line 6, in <lambda>(_, args)
----> 6     lambda _, args: MyTree(*args))

Cell In[10], line 3, in MyTree.__init__(self, a)
      2   def __init__(self, a):
----> 3     self.a = jnp.asarray(a)

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/array_constructors.py:428, in asarray(a, dtype, order, copy, device, out_sharding)
    426 if dtype is not None:
    427   dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "asarray")
--> 428 return array(a, dtype=dtype, copy=bool(copy), order=order, device=device,
    429              out_sharding=out_sharding)

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/array_constructors.py:260, in array(object, dtype, copy, order, ndmin, device, out_sharding)
    257 leaves, treedef = tree_util.tree_flatten(
    258     object, is_leaf=lambda x: not isinstance(x, (list, tuple)))
    259 if any(leaf is None for leaf in leaves):
--> 260   raise ValueError("None is not a valid value for jnp.array")
    261 leaves = [
    262     leaf
    263     if (leaf_jax_array := getattr(leaf, "__jax_array__", None)) is None
    264     else leaf_jax_array()
    265     for leaf in leaves
    266 ]
    267 if dtype is None:
    268   # Use lattice_result_type rather than result_type to avoid canonicalization.
    269   # Otherwise, weakly-typed inputs would have their dtypes canonicalized.

ValueError: None is not a valid value for jnp.array

• 在第一个使用 jax.vmap(...)(tree) 的例子中，JAX 内部使用 object() 值的数组来推断树的结构。

• 在第二个使用 jax.jacobian(...)(tree) 的例子中，将树映射到树的函数的雅可比矩阵被定义为树的树。

解决方案 1

• 自定义 pytree 类的 __init__ 和 __new__ 方法通常应避免进行任何数组转换或其他输入验证，否则就需要预见并处理这些特殊情况。例如：

class MyTree:
  def __init__(self, a):
    if not (type(a) is object or a is None or isinstance(a, MyTree)):
      a = jnp.asarray(a)
    self.a = a

解决方案 2

• 构建你的自定义 tree_unflatten 函数，使其避免调用 __init__。如果你选择这条路，请确保当代码更新时，你的 tree_unflatten 函数与 __init__ 保持同步。示例：

def tree_unflatten(aux_data, children):
  del aux_data  # Unused in this class.
  obj = object.__new__(MyTree)
  obj.a = children[0]
  return obj

内部 pytree 处理

JAX 在 api.py 边界（以及控制流原语中）将 pytree 展平为叶子列表。这保持了下游 JAX 内部结构的简洁：像 grad()、jit() 和 vmap() 这样的转换可以处理接收和返回各种不同 Python 容器的用户函数，而系统的所有其他部分可以仅对接收（多个）数组参数并始终返回数组平坦列表的函数进行操作。

当 JAX 展平一个 pytree 时，它会生成一个叶子列表和一个编码原始值结构的 treedef 对象。然后，在对叶子进行转换后，该 treedef 可用于构建匹配的结构化值。Pytree 是树状的，而不是 DAG 或图状的，因为我们假设它们具有引用透明性，并且不包含引用循环。

这是一个简单的例子：

from jax.tree_util import tree_flatten, tree_unflatten
import jax.numpy as jnp

# The structured value to be transformed
value_structured = [1., (2., 3.)]

# The leaves in value_flat correspond to the `*` markers in value_tree
value_flat, value_tree = tree_flatten(value_structured)
print(f"{value_flat=}\n{value_tree=}")

# Transform the flat value list using an element-wise numeric transformer
transformed_flat = list(map(lambda v: v * 2., value_flat))
print(f"{transformed_flat=}")

# Reconstruct the structured output, using the original
transformed_structured = tree_unflatten(value_tree, transformed_flat)
print(f"{transformed_structured=}")

value_flat=[1.0, 2.0, 3.0]
value_tree=PyTreeDef([*, (*, *)])
transformed_flat=[2.0, 4.0, 6.0]
transformed_structured=[2.0, (4.0, 6.0)]

默认情况下，pytree 容器可以是列表、元组、字典、namedtuple、None、OrderedDict。其他类型的值（包括数字和 ndarray 值）被视为叶子。

from collections import namedtuple
Point = namedtuple('Point', ['x', 'y'])

example_containers = [
    (1., [2., 3.]),
    (1., {'b': 2., 'a': 3.}),
    1.,
    None,
    jnp.zeros(2),
    Point(1., 2.)
]
def show_example(structured):
  flat, tree = tree_flatten(structured)
  unflattened = tree_unflatten(tree, flat)
  print(f"{structured=}\n  {flat=}\n  {tree=}\n  {unflattened=}")

for structured in example_containers:
  show_example(structured)

structured=(1.0, [2.0, 3.0])
  flat=[1.0, 2.0, 3.0]
  tree=PyTreeDef((*, [*, *]))
  unflattened=(1.0, [2.0, 3.0])
structured=(1.0, {'b': 2.0, 'a': 3.0})
  flat=[1.0, 3.0, 2.0]
  tree=PyTreeDef((*, {'a': *, 'b': *}))
  unflattened=(1.0, {'a': 3.0, 'b': 2.0})
structured=1.0
  flat=[1.0]
  tree=PyTreeDef(*)
  unflattened=1.0
structured=None
  flat=[]
  tree=PyTreeDef(None)
  unflattened=None
structured=Array([0., 0.], dtype=float32)
  flat=[Array([0., 0.], dtype=float32)]
  tree=PyTreeDef(*)
  unflattened=Array([0., 0.], dtype=float32)
structured=Point(x=1.0, y=2.0)
  flat=[1.0, 2.0]
  tree=PyTreeDef(CustomNode(namedtuple[Point], [*, *]))
  unflattened=Point(x=1.0, y=2.0)