jax.tree_util.register_pytree_with_keys

jax.tree_util.register_pytree_with_keys#

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

扩展了在 pytrees 中被视为内部节点的类型集合。

这是一个比 register_pytree_node 更强大的替代方案，允许你在展平和树映射时访问每个 pytree 叶子的键路径。

参数:

  • nodetype (type[T]) – 一个被视为内部 pytree 节点的 Python 类型。

  • flatten_with_keys (Callable[[T], tuple[Iterable[tuple[KeyEntry, Any]], _AuxData]]) – 在扁平化过程中使用的一个函数，接受类型为 nodetype 的值并返回一对，其中 (1) 是一个包含每个键路径及其子项的元组的可迭代对象，(2) 是一些可哈希的辅助数据，这些数据将被存储在 treedef 中并传递给 unflatten_func。

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

  • flatten_func (None | Callable[[T], tuple[Iterable[Any], _AuxData]]) – 一个类似于 flatten_with_keys 的可选函数，但只返回子节点和辅助数据。它必须以与 flatten_with_keys 相同的顺序返回子节点，并返回相同的辅助数据。此参数是可选的，仅在调用没有键的函数如 tree_map 和 tree_flatten 时需要更快的遍历时才需要。

示例

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

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

现在使用一个键感知的扁平化函数来注册它：

>>> from jax.tree_util import register_pytree_with_keys_class, GetAttrKey
>>> def flatten_with_keys(obj):
...   children = [(GetAttrKey('x'), obj.x),
...               (GetAttrKey('y'), obj.y)]  # children must contain arrays & pytrees
...   aux_data = (obj.size,)  # aux_data must contain static, hashable data.
...   return children, aux_data
...
>>> def unflatten(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_with_keys, unflatten)

现在这可以与 tree_flatten_with_path() 这样的函数一起使用：

>>> m = MyContainer(4)
>>> leaves, treedef = jax.tree_util.tree_flatten_with_path(m)
目录