jax.tree_util.register_pytree_node

jax.tree_util.register_pytree_node#

jax.tree_util.register_pytree_node(nodetype, flatten_func, unflatten_func)[source]#

扩展被认为是 pytrees 中内部节点的一组类型。

参见 示例用法.

参数:

  • 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 的实例。

返回类型:

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)
内容