jax.tree_util.register_pytree_node#

jax.tree_util.register_pytree_node(nodetype, flatten_func, unflatten_func, flatten_with_keys_func=None)[源代码]#

扩展了在 pytree 中被认为是内部节点的类型集合。

请参阅 用法示例。

参数:

  • nodetype (type[T]) – 要注册为 pytree 的 Python 类型。

  • flatten_func (Callable[[T], tuple[_Children, _AuxData]]) – 一个在展平过程中使用的函数，它接受一个类型为 nodetype 的值，并返回一个配对，其中 (1) 是一个可迭代对象，用于递归展平子节点，(2) 是一个可哈希的辅助数据，用于存储在 treedef 中，并传递给 unflatten_func。

  • unflatten_func (Callable[[_AuxData, _Children], T]) – 一个接受两个参数的函数：由 flatten_func 返回并存储在 treedef 中的辅助数据，以及未展平的子节点。该函数应返回一个 nodetype 的实例。

  • flatten_with_keys_func (Callable[[T], tuple[KeyLeafPairs, _AuxData]] | None | None)

返回类型:

None

另请参阅

示例

首先，我们将定义一个自定义类型

>>> class MyContainer:
...   def __init__(self, size):
...     self.x = jnp.zeros(size)
...     self.y = jnp.ones(size)
...     self.size = size

如果我们尝试在 JIT 编译的函数中使用此类型，我们会收到错误，因为 JAX 尚不知道如何处理此类型

>>> m = MyContainer(size=5)
>>> def f(m):
...   return m.x + m.y + jnp.arange(m.size)
>>> jax.jit(f)(m)  
Traceback (most recent call last):
  ...
TypeError: Cannot interpret value of type <class 'jax.tree_util.MyContainer'> as an abstract array; it does not have a dtype attribute

为了使我们的对象被 JAX 识别，我们必须将其注册为 pytree

>>> def flatten_func(obj):
...   children = (obj.x, obj.y)  # children must contain arrays & pytrees
...   aux_data = (obj.size,)  # aux_data must contain static, hashable data.
...   return (children, aux_data)
...
>>> def unflatten_func(aux_data, children):
...   # Here we avoid `__init__` because it has extra logic we don't require:
...   obj = object.__new__(MyContainer)
...   obj.x, obj.y = children
...   obj.size, = aux_data
...   return obj
...
>>> jax.tree_util.register_pytree_node(MyContainer, flatten_func, unflatten_func)

现在定义好之后，我们可以在 JIT 编译的函数中使用此类型的实例。

>>> jax.jit(f)(m)
Array([1., 2., 3., 4., 5.], dtype=float32)