Tree Operations

Constants

`MAX_RECURSION_DEPTH`	Maximum recursion depth for pytree traversal.
`NONE_IS_NODE`	Literal constant that treats `None` as a pytree non-leaf node.
`NONE_IS_LEAF`	Literal constant that treats `None` as a pytree leaf node.

optree.MAX_RECURSION_DEPTH: int = 1000

Maximum recursion depth for pytree traversal. It is 1000.

This limit prevents infinite recursion from causing an overflow of the C stack and crashing Python.

optree.NONE_IS_NODE: bool = False: Literal constant that treats None as a pytree non-leaf node.

optree.NONE_IS_LEAF: bool = True: Literal constant that treats None as a pytree leaf node.

Tree Manipulation Functions

`tree_flatten`(tree[, is_leaf, none_is_leaf, ...])	Flatten a pytree.
`tree_flatten_with_path`(tree[, is_leaf, ...])	Flatten a pytree and additionally record the paths.
`tree_flatten_with_accessor`(tree[, is_leaf, ...])	Flatten a pytree and additionally record the accessors.
`tree_unflatten`(treespec, leaves)	Reconstruct a pytree from the treespec and the leaves.
`tree_iter`(tree[, is_leaf, none_is_leaf, ...])	Get an iterator over the leaves of a pytree.
`tree_leaves`(tree[, is_leaf, none_is_leaf, ...])	Get the leaves of a pytree.
`tree_structure`(tree[, is_leaf, ...])	Get the treespec for a pytree.
`tree_paths`(tree[, is_leaf, none_is_leaf, ...])	Get the path entries to the leaves of a pytree.
`tree_accessors`(tree[, is_leaf, ...])	Get the accessors to the leaves of a pytree.
`tree_is_leaf`(tree[, is_leaf, none_is_leaf, ...])	Test whether the given object is a leaf node.
`all_leaves`(iterable[, is_leaf, ...])	Test whether all elements in the given iterable are all leaves.
`tree_map`(func, tree, *rests[, is_leaf, ...])	Map a multi-input function over pytree args to produce a new pytree.
`tree_map_`(func, tree, *rests[, is_leaf, ...])	Like `tree_map()`, but do an inplace call on each leaf and return the original tree.
`tree_map_with_path`(func, tree, *rests[, ...])	Map a multi-input function over pytree args as well as the tree paths to produce a new pytree.
`tree_map_with_path_`(func, tree, *rests[, ...])	Like `tree_map_with_path()`, but do an inplace call on each leaf and return the original tree.
`tree_map_with_accessor`(func, tree, *rests[, ...])	Map a multi-input function over pytree args as well as the tree accessors to produce a new pytree.
`tree_map_with_accessor_`(func, tree, *rests)	Like `tree_map_with_accessor()`, but do an inplace call on each leaf and return the original tree.
`tree_replace_nones`(sentinel, tree[, namespace])	Replace `None` in `tree` with `sentinel`.
`tree_transpose`(outer_treespec, ...[, is_leaf])	Transform a tree having tree structure (outer, inner) into one having structure (inner, outer).
`tree_transpose_map`(func, tree, *rests[, ...])	Map a multi-input function over pytree args to produce a new pytree with transposed structure.
`tree_transpose_map_with_path`(func, tree, *rests)	Map a multi-input function over pytree args as well as the tree paths to produce a new pytree with transposed structure.
`tree_transpose_map_with_accessor`(func, tree, ...)	Map a multi-input function over pytree args as well as the tree accessors to produce a new pytree with transposed structure.
`tree_broadcast_prefix`(prefix_tree, full_tree)	Return a pytree of same structure of `full_tree` with broadcasted subtrees in `prefix_tree`.
`broadcast_prefix`(prefix_tree, full_tree[, ...])	Return a list of broadcasted leaves in `prefix_tree` to match the number of leaves in `full_tree`.
`tree_broadcast_common`(tree, other_tree[, ...])	Return two pytrees of common suffix structure of `tree` and `other_tree` with broadcasted subtrees.
`broadcast_common`(tree, other_tree[, ...])	Return two lists of leaves in `tree` and `other_tree` broadcasted to match the number of leaves in the common suffix structure.
`tree_broadcast_map`(func, tree, *rests[, ...])	Map a multi-input function over pytree args to produce a new pytree.
`tree_broadcast_map_with_path`(func, tree, *rests)	Map a multi-input function over pytree args as well as the tree paths to produce a new pytree.
`tree_broadcast_map_with_accessor`(func, tree, ...)	Map a multi-input function over pytree args as well as the tree accessors to produce a new pytree.
`tree_flatten_one_level`(tree[, is_leaf, ...])	Flatten the pytree one level, returning a 4-tuple of children, auxiliary data, path entries, and an unflatten function.
`prefix_errors`(prefix_tree, full_tree[, ...])	Return a list of errors that would be raised by `broadcast_prefix()`.

optree.tree_flatten(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Flatten a pytree.

The flattening order (i.e., the order of elements in the output list) is deterministic, corresponding to a left-to-right depth-first tree traversal.

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}
>>> tree_flatten(tree)  
(
    [1, 2, 3, 4, 5],
    PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': None, 'd': *})
)
>>> tree_flatten(tree, none_is_leaf=True)  
(
    [1, 2, 3, 4, None, 5],
    PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': *, 'd': *}, NoneIsLeaf)
)
>>> tree_flatten(1)
([1], PyTreeSpec(*))
>>> tree_flatten(None)
([], PyTreeSpec(None))
>>> tree_flatten(None, none_is_leaf=True)
([None], PyTreeSpec(*, NoneIsLeaf))

For unordered dictionaries, dict and collections.defaultdict, the order is dependent on the sorted keys in the dictionary. Please use collections.OrderedDict if you want to keep the keys in the insertion order.

>>> from collections import OrderedDict
>>> tree = OrderedDict([('b', (2, [3, 4])), ('a', 1), ('c', None), ('d', 5)])
>>> tree_flatten(tree)  
(
    [2, 3, 4, 1, 5],
    PyTreeSpec(OrderedDict({'b': (*, [*, *]), 'a': *, 'c': None, 'd': *}))
)
>>> tree_flatten(tree, none_is_leaf=True)  
(
    [2, 3, 4, 1, None, 5],
    PyTreeSpec(OrderedDict({'b': (*, [*, *]), 'a': *, 'c': *, 'd': *}), NoneIsLeaf)
)

Parameters:

tree (pytree) – A pytree to flatten.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

tuple[list[TypeVar(T)], PyTreeSpec]

Returns:

A pair (leaves, treespec) where the first element is a list of leaf values and the second element is a treespec representing the structure of the pytree.

optree.tree_flatten_with_path(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Flatten a pytree and additionally record the paths.

The flattening order (i.e., the order of elements in the output list) is deterministic, corresponding to a left-to-right depth-first tree traversal.

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}
>>> tree_flatten_with_path(tree)  
(
    [('a',), ('b', 0), ('b', 1, 0), ('b', 1, 1), ('d',)],
    [1, 2, 3, 4, 5],
    PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': None, 'd': *})
)
>>> tree_flatten_with_path(tree, none_is_leaf=True)  
(
    [('a',), ('b', 0), ('b', 1, 0), ('b', 1, 1), ('c',), ('d',)],
    [1, 2, 3, 4, None, 5],
    PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': *, 'd': *}, NoneIsLeaf)
)
>>> tree_flatten_with_path(1)
([()], [1], PyTreeSpec(*))
>>> tree_flatten_with_path(None)
([], [], PyTreeSpec(None))
>>> tree_flatten_with_path(None, none_is_leaf=True)
([()], [None], PyTreeSpec(*, NoneIsLeaf))

For unordered dictionaries, dict and collections.defaultdict, the order is dependent on the sorted keys in the dictionary. Please use collections.OrderedDict if you want to keep the keys in the insertion order.

>>> from collections import OrderedDict
>>> tree = OrderedDict([('b', (2, [3, 4])), ('a', 1), ('c', None), ('d', 5)])
>>> tree_flatten_with_path(tree)  
(
    [('b', 0), ('b', 1, 0), ('b', 1, 1), ('a',), ('d',)],
    [2, 3, 4, 1, 5],
    PyTreeSpec(OrderedDict({'b': (*, [*, *]), 'a': *, 'c': None, 'd': *}))
)
>>> tree_flatten_with_path(tree, none_is_leaf=True)  
(
    [('b', 0), ('b', 1, 0), ('b', 1, 1), ('a',), ('c',), ('d',)],
    [2, 3, 4, 1, None, 5],
    PyTreeSpec(OrderedDict({'b': (*, [*, *]), 'a': *, 'c': *, 'd': *}), NoneIsLeaf)
)

Parameters:

tree (pytree) – A pytree to flatten.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

tuple[list[tuple[Any, ...]], list[TypeVar(T)], PyTreeSpec]

Returns:

A triple (paths, leaves, treespec). The first element is a list of the paths to the leaf values, while each path is a tuple of the index or keys. The second element is a list of leaf values and the last element is a treespec representing the structure of the pytree.

optree.tree_flatten_with_accessor(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Flatten a pytree and additionally record the accessors.

The flattening order (i.e., the order of elements in the output list) is deterministic, corresponding to a left-to-right depth-first tree traversal.

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}
>>> tree_flatten_with_accessor(tree)  
(
    [
        PyTreeAccessor(*['a'], (MappingEntry(key='a', type=<class 'dict'>),)),
        PyTreeAccessor(*['b'][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
        PyTreeAccessor(*['b'][1][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=0, type=<class 'list'>))),
        PyTreeAccessor(*['b'][1][1], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=1, type=<class 'list'>))),
        PyTreeAccessor(*['d'], (MappingEntry(key='d', type=<class 'dict'>),))
    ],
    [1, 2, 3, 4, 5],
    PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': None, 'd': *})
)
>>> tree_flatten_with_accessor(tree, none_is_leaf=True)  
(
    [
        PyTreeAccessor(*['a'], (MappingEntry(key='a', type=<class 'dict'>),)),
        PyTreeAccessor(*['b'][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
        PyTreeAccessor(*['b'][1][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=0, type=<class 'list'>))),
        PyTreeAccessor(*['b'][1][1], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=1, type=<class 'list'>))),
        PyTreeAccessor(*['c'], (MappingEntry(key='c', type=<class 'dict'>),)),
        PyTreeAccessor(*['d'], (MappingEntry(key='d', type=<class 'dict'>),))
    ],
    [1, 2, 3, 4, None, 5],
    PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': *, 'd': *}, NoneIsLeaf)
)
>>> tree_flatten_with_accessor(1)
([PyTreeAccessor(*, ())], [1], PyTreeSpec(*))
>>> tree_flatten_with_accessor(None)
([], [], PyTreeSpec(None))
>>> tree_flatten_with_accessor(None, none_is_leaf=True)
([PyTreeAccessor(*, ())], [None], PyTreeSpec(*, NoneIsLeaf))

For unordered dictionaries, dict and collections.defaultdict, the order is dependent on the sorted keys in the dictionary. Please use collections.OrderedDict if you want to keep the keys in the insertion order.

>>> from collections import OrderedDict
>>> tree = OrderedDict([('b', (2, [3, 4])), ('a', 1), ('c', None), ('d', 5)])
>>> tree_flatten_with_accessor(tree)  
(
    [
        PyTreeAccessor(*['b'][0], (MappingEntry(key='b', type=<class 'collections.OrderedDict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
        PyTreeAccessor(*['b'][1][0], (MappingEntry(key='b', type=<class 'collections.OrderedDict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=0, type=<class 'list'>))),
        PyTreeAccessor(*['b'][1][1], (MappingEntry(key='b', type=<class 'collections.OrderedDict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=1, type=<class 'list'>))),
        PyTreeAccessor(*['a'], (MappingEntry(key='a', type=<class 'collections.OrderedDict'>),)),
        PyTreeAccessor(*['d'], (MappingEntry(key='d', type=<class 'collections.OrderedDict'>),))
    ],
    [2, 3, 4, 1, 5],
    PyTreeSpec(OrderedDict({'b': (*, [*, *]), 'a': *, 'c': None, 'd': *}))
)
>>> tree_flatten_with_accessor(tree, none_is_leaf=True)  
(
    [
        PyTreeAccessor(*['b'][0], (MappingEntry(key='b', type=<class 'collections.OrderedDict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
        PyTreeAccessor(*['b'][1][0], (MappingEntry(key='b', type=<class 'collections.OrderedDict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=0, type=<class 'list'>))),
        PyTreeAccessor(*['b'][1][1], (MappingEntry(key='b', type=<class 'collections.OrderedDict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=1, type=<class 'list'>))),
        PyTreeAccessor(*['a'], (MappingEntry(key='a', type=<class 'collections.OrderedDict'>),)),
        PyTreeAccessor(*['c'], (MappingEntry(key='c', type=<class 'collections.OrderedDict'>),)),
        PyTreeAccessor(*['d'], (MappingEntry(key='d', type=<class 'collections.OrderedDict'>),))
    ],
    [2, 3, 4, 1, None, 5],
    PyTreeSpec(OrderedDict({'b': (*, [*, *]), 'a': *, 'c': *, 'd': *}), NoneIsLeaf)
)

Parameters:

tree (pytree) – A pytree to flatten.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

tuple[list[PyTreeAccessor], list[TypeVar(T)], PyTreeSpec]

Returns:

A triple (accessors, leaves, treespec). The first element is a list of accessors to the leaf values. The second element is a list of leaf values and the last element is a treespec representing the structure of the pytree.

optree.tree_unflatten(treespec, leaves)[source]

Reconstruct a pytree from the treespec and the leaves.

The inverse of tree_flatten().

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}
>>> leaves, treespec = tree_flatten(tree)
>>> tree == tree_unflatten(treespec, leaves)
True

Parameters:

treespec (PyTreeSpec) – The treespec to reconstruct.
leaves (iterable) – The list of leaves to use for reconstruction. The list must match the number of leaves of the treespec.

Return type:

Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]

Returns:

The reconstructed pytree, containing the leaves placed in the structure described by treespec.

optree.tree_iter(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Get an iterator over the leaves of a pytree.

See also tree_flatten() and tree_leaves().

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}
>>> list(tree_iter(tree))
[1, 2, 3, 4, 5]
>>> list(tree_iter(tree, none_is_leaf=True))
[1, 2, 3, 4, None, 5]
>>> list(tree_iter(1))
[1]
>>> list(tree_iter(None))
[]
>>> list(tree_iter(None, none_is_leaf=True))
[None]

Parameters:

tree (pytree) – A pytree to iterate over.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Iterable[TypeVar(T)]

Returns:

An iterator over the leaf values.

optree.tree_leaves(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Get the leaves of a pytree.

See also tree_flatten() and tree_iter().

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}
>>> tree_leaves(tree)
[1, 2, 3, 4, 5]
>>> tree_leaves(tree, none_is_leaf=True)
[1, 2, 3, 4, None, 5]
>>> tree_leaves(1)
[1]
>>> tree_leaves(None)
[]
>>> tree_leaves(None, none_is_leaf=True)
[None]

Parameters:

tree (pytree) – A pytree to flatten.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

list[TypeVar(T)]

Returns:

A list of leaf values.

optree.tree_structure(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Get the treespec for a pytree.

See also tree_flatten().

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}
>>> tree_structure(tree)
PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': None, 'd': *})
>>> tree_structure(tree, none_is_leaf=True)
PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': *, 'd': *}, NoneIsLeaf)
>>> tree_structure(1)
PyTreeSpec(*)
>>> tree_structure(None)
PyTreeSpec(None)
>>> tree_structure(None, none_is_leaf=True)
PyTreeSpec(*, NoneIsLeaf)

Parameters:

tree (pytree) – A pytree to flatten.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec object representing the structure of the pytree.

optree.tree_paths(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Get the path entries to the leaves of a pytree.

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}
>>> tree_paths(tree)
[('a',), ('b', 0), ('b', 1, 0), ('b', 1, 1), ('d',)]
>>> tree_paths(tree, none_is_leaf=True)
[('a',), ('b', 0), ('b', 1, 0), ('b', 1, 1), ('c',), ('d',)]
>>> tree_paths(1)
[()]
>>> tree_paths(None)
[]
>>> tree_paths(None, none_is_leaf=True)
[()]

Parameters:

tree (pytree) – A pytree to flatten.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

list[tuple[Any, ...]]

Returns:

A list of the paths to the leaf values, while each path is a tuple of the index or keys.

optree.tree_accessors(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Get the accessors to the leaves of a pytree.

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}
>>> tree_accessors(tree)  
[
    PyTreeAccessor(*['a'], (MappingEntry(key='a', type=<class 'dict'>),)),
    PyTreeAccessor(*['b'][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
    PyTreeAccessor(*['b'][1][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=0, type=<class 'list'>))),
    PyTreeAccessor(*['b'][1][1], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=1, type=<class 'list'>))),
    PyTreeAccessor(*['d'], (MappingEntry(key='d', type=<class 'dict'>),))
]
>>> tree_accessors(tree, none_is_leaf=True)  
[
    PyTreeAccessor(*['a'], (MappingEntry(key='a', type=<class 'dict'>),)),
    PyTreeAccessor(*['b'][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
    PyTreeAccessor(*['b'][1][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=0, type=<class 'list'>))),
    PyTreeAccessor(*['b'][1][1], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=1, type=<class 'list'>))),
    PyTreeAccessor(*['c'], (MappingEntry(key='c', type=<class 'dict'>),)),
    PyTreeAccessor(*['d'], (MappingEntry(key='d', type=<class 'dict'>),))
]
>>> tree_accessors(1)
[PyTreeAccessor(*, ())]
>>> tree_accessors(None)
[]
>>> tree_accessors(None, none_is_leaf=True)
[PyTreeAccessor(*, ())]

Parameters:

tree (pytree) – A pytree to flatten.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

list[PyTreeAccessor]

Returns:

A list of accessors to the leaf values.

optree.tree_is_leaf(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Test whether the given object is a leaf node.

>>> tree_is_leaf(1)
True
>>> tree_is_leaf(None)
False
>>> tree_is_leaf(None, none_is_leaf=True)
True
>>> tree_is_leaf({'a': 1, 'b': (2, 3)})
False

Parameters:

tree (pytree) – A pytree to check if it is a leaf node.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than a leaf. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

bool

Returns:

A boolean indicating if all elements in the input iterable are leaves.

optree.all_leaves(iterable, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Test whether all elements in the given iterable are all leaves.

>>> tree = {'a': [1, 2, 3]}
>>> all_leaves(tree_leaves(tree))
True
>>> all_leaves([tree])
False
>>> all_leaves([1, 2, None, 3])
False
>>> all_leaves([1, 2, None, 3], none_is_leaf=True)
True

Note that this function iterates and checks the elements in the input iterable object, which uses the iter() function. For dictionaries, iter(d) for a dictionary d iterates the keys of the dictionary, not the values.

>>> list({'a': 1, 'b': (2, 3)})
['a', 'b']
>>> all_leaves({'a': 1, 'b': (2, 3)})
True

This function is useful in advanced cases. For example, if a library allows arbitrary map operations on a flat list of leaves it may want to check if the result is still a flat list of leaves.

Parameters:

iterable (iterable) – A iterable of leaves.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than a leaf. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

bool

Returns:

A boolean indicating if all elements in the input iterable are leaves.

optree.tree_map(func, tree, *rests, is_leaf=None, none_is_leaf=False, namespace='')[source]

Map a multi-input function over pytree args to produce a new pytree.

>>> tree_map(lambda x: x + 1, {'x': 7, 'y': (42, 64)})
{'x': 8, 'y': (43, 65)}
>>> tree_map(lambda x: x + 1, {'x': 7, 'y': (42, 64), 'z': None})
{'x': 8, 'y': (43, 65), 'z': None}
>>> tree_map(lambda x: x is None, {'x': 7, 'y': (42, 64), 'z': None})
{'x': False, 'y': (False, False), 'z': None}
>>> tree_map(lambda x: x is None, {'x': 7, 'y': (42, 64), 'z': None}, none_is_leaf=True)
{'x': False, 'y': (False, False), 'z': True}

If multiple inputs are given, the structure of the tree is taken from the first input; subsequent inputs need only have tree as a prefix:

>>> tree_map(lambda x, y: [x] + y, [5, 6], [[7, 9], [1, 2]])
[[5, 7, 9], [6, 1, 2]]

Parameters:

func (callable) – A function that takes 1 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees.
tree (pytree) – A pytree to be mapped over, with each leaf providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, each of which has the same structure as tree or has tree as a prefix.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(U), Tuple[PyTree[U], ...], List[PyTree[U]], Dict[Any, PyTree[U]], Deque[PyTree[U]], CustomTreeNode[PyTree[U]]]

Returns:

A new pytree with the same structure as tree but with the value at each leaf given by func(x, *xs) where x is the value at the corresponding leaf in tree and xs is the tuple of values at corresponding nodes in rests.

optree.tree_map_(func, tree, *rests, is_leaf=None, none_is_leaf=False, namespace='')[source]

Like tree_map(), but do an inplace call on each leaf and return the original tree.

Parameters:

func (callable) – A function that takes 1 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees.
tree (pytree) – A pytree to be mapped over, with each leaf providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, each of which has the same structure as tree or has tree as a prefix.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]

Returns:

The original tree with the value at each leaf is given by the side-effect of function func(x, *xs) (not the return value) where x is the value at the corresponding leaf in tree and xs is the tuple of values at values at corresponding nodes in rests.

optree.tree_map_with_path(func, tree, *rests, is_leaf=None, none_is_leaf=False, namespace='')[source]

Map a multi-input function over pytree args as well as the tree paths to produce a new pytree.

>>> tree_map_with_path(lambda p, x: (len(p), x), {'x': 7, 'y': (42, 64)})
{'x': (1, 7), 'y': ((2, 42), (2, 64))}
>>> tree_map_with_path(lambda p, x: x + len(p), {'x': 7, 'y': (42, 64), 'z': None})
{'x': 8, 'y': (44, 66), 'z': None}
>>> tree_map_with_path(lambda p, x: p, {'x': 7, 'y': (42, 64), 'z': {1.5: None}})
{'x': ('x',), 'y': (('y', 0), ('y', 1)), 'z': {1.5: None}}
>>> tree_map_with_path(lambda p, x: p, {'x': 7, 'y': (42, 64), 'z': {1.5: None}}, none_is_leaf=True)
{'x': ('x',), 'y': (('y', 0), ('y', 1)), 'z': {1.5: ('z', 1.5)}}

Parameters:

func (callable) – A function that takes 2 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees with extra paths.
tree (pytree) – A pytree to be mapped over, with each leaf providing the second positional argument and the corresponding path providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, each of which has the same structure as tree or has tree as a prefix.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(U), Tuple[PyTree[U], ...], List[PyTree[U]], Dict[Any, PyTree[U]], Deque[PyTree[U]], CustomTreeNode[PyTree[U]]]

Returns:

A new pytree with the same structure as tree but with the value at each leaf given by func(p, x, *xs) where (p, x) are the path and value at the corresponding leaf in tree and xs is the tuple of values at corresponding nodes in rests.

optree.tree_map_with_path_(func, tree, *rests, is_leaf=None, none_is_leaf=False, namespace='')[source]

Like tree_map_with_path(), but do an inplace call on each leaf and return the original tree.

Parameters:

func (callable) – A function that takes 2 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees with extra paths.
tree (pytree) – A pytree to be mapped over, with each leaf providing the second positional argument and the corresponding path providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, each of which has the same structure as tree or has tree as a prefix.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]

Returns:

The original tree with the value at each leaf is given by the side-effect of function func(p, x, *xs) (not the return value) where (p, x) are the path and value at the corresponding leaf in tree and xs is the tuple of values at values at corresponding nodes in rests.

optree.tree_map_with_accessor(func, tree, *rests, is_leaf=None, none_is_leaf=False, namespace='')[source]

Map a multi-input function over pytree args as well as the tree accessors to produce a new pytree.

>>> tree_map_with_accessor(lambda a, x: f'{a.codegen("tree")} = {x!r}', {'x': 7, 'y': (42, 64)})
{'x': "tree['x'] = 7", 'y': ("tree['y'][0] = 42", "tree['y'][1] = 64")}
>>> tree_map_with_accessor(lambda a, x: x + len(a), {'x': 7, 'y': (42, 64), 'z': None})
{'x': 8, 'y': (44, 66), 'z': None}
>>> tree_map_with_accessor(  
...     lambda a, x: a,
...     {'x': 7, 'y': (42, 64), 'z': {1.5: None}},
... )
{
    'x': PyTreeAccessor(*['x'], (MappingEntry(key='x', type=<class 'dict'>),)),
    'y': (
        PyTreeAccessor(*['y'][0], (MappingEntry(key='y', type=<class 'dict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
        PyTreeAccessor(*['y'][1], (MappingEntry(key='y', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>)))
    ),
    'z': {1.5: None}
}
>>> tree_map_with_accessor(  
...     lambda a, x: a,
...     {'x': 7, 'y': (42, 64), 'z': {1.5: None}},
...     none_is_leaf=True,
... )
{
    'x': PyTreeAccessor(*['x'], (MappingEntry(key='x', type=<class 'dict'>),)),
    'y': (
        PyTreeAccessor(*['y'][0], (MappingEntry(key='y', type=<class 'dict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
        PyTreeAccessor(*['y'][1], (MappingEntry(key='y', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>)))
    ),
    'z': {
        1.5: PyTreeAccessor(*['z'][1.5], (MappingEntry(key='z', type=<class 'dict'>), MappingEntry(key=1.5, type=<class 'dict'>)))
    }
}

Parameters:

func (callable) – A function that takes 2 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees with extra accessors.
tree (pytree) – A pytree to be mapped over, with each leaf providing the second positional argument and the corresponding path providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, each of which has the same structure as tree or has tree as a prefix.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(U), Tuple[PyTree[U], ...], List[PyTree[U]], Dict[Any, PyTree[U]], Deque[PyTree[U]], CustomTreeNode[PyTree[U]]]

Returns:

A new pytree with the same structure as tree but with the value at each leaf given by func(a, x, *xs) where (a, x) are the accessor and value at the corresponding leaf in tree and xs is the tuple of values at corresponding nodes in rests.

optree.tree_map_with_accessor_(func, tree, *rests, is_leaf=None, none_is_leaf=False, namespace='')[source]

Like tree_map_with_accessor(), but do an inplace call on each leaf and return the original tree.

Parameters:

func (callable) – A function that takes 2 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees with extra accessors.
tree (pytree) – A pytree to be mapped over, with each leaf providing the second positional argument and the corresponding path providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, each of which has the same structure as tree or has tree as a prefix.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]

Returns:

The original tree with the value at each leaf is given by the side-effect of function func(a, x, *xs) (not the return value) where (a, x) are the accessor and value at the corresponding leaf in tree and xs is the tuple of values at values at corresponding nodes in rests.

optree.tree_replace_nones(sentinel, tree, namespace='')[source]

Replace None in tree with sentinel.

See also tree_flatten() and tree_map().

>>> tree_replace_nones(0, {'a': 1, 'b': None, 'c': (2, None)})
{'a': 1, 'b': 0, 'c': (2, 0)}
>>> tree_replace_nones(0, None)
0

Parameters:

sentinel (object) – The value to replace None with.
tree (pytree) – A pytree to be transformed.
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]

Returns:

A new pytree with the same structure as tree but with None replaced.

optree.tree_transpose(outer_treespec, inner_treespec, tree, is_leaf=None)[source]

Transform a tree having tree structure (outer, inner) into one having structure (inner, outer).

>>> outer_treespec = tree_structure({'a': 1, 'b': 2, 'c': (3, 4)})
>>> outer_treespec
PyTreeSpec({'a': *, 'b': *, 'c': (*, *)})
>>> inner_treespec = tree_structure((1, 2))
>>> inner_treespec
PyTreeSpec((*, *))
>>> tree = {'a': (1, 2), 'b': (3, 4), 'c': ((5, 6), (7, 8))}
>>> tree_transpose(outer_treespec, inner_treespec, tree)
({'a': 1, 'b': 3, 'c': (5, 7)}, {'a': 2, 'b': 4, 'c': (6, 8)})

For performance reasons, this function is only checks for the number of leaves in the input pytree, not the structure. The result is only enumerated up to the original order of leaves in tree, then transpose depends on the number of leaves in structure (inner, outer). The caller is responsible for ensuring that the input pytree has a prefix structure of outer_treespec followed by a prefix structure of inner_treespec. Otherwise, the result may be incorrect.

>>> tree_transpose(outer_treespec, inner_treespec, list(range(1, 9)))
({'a': 1, 'b': 3, 'c': (5, 7)}, {'a': 2, 'b': 4, 'c': (6, 8)})

Parameters:

outer_treespec (PyTreeSpec) – A treespec object representing the outer structure of the pytree.
inner_treespec (PyTreeSpec) – A treespec object representing the inner structure of the pytree.
tree (pytree) – A pytree to be transposed.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.

Return type:

Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]

Returns:

A new pytree with the same structure as inner_treespec but with the value at each leaf has the same structure as outer_treespec.

optree.tree_transpose_map(func, tree, *rests, inner_treespec=None, is_leaf=None, none_is_leaf=False, namespace='')[source]

Map a multi-input function over pytree args to produce a new pytree with transposed structure.

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)}
>>> tree_transpose_map(  
...     lambda x: {'identity': x, 'double': 2 * x},
...     tree,
... )
{
    'identity': {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)},
    'double': {'b': (4, [6, 8]), 'a': 2, 'c': (10, 12)}
}
>>> tree_transpose_map(  
...     lambda x: {'identity': x, 'double': (x, x)},
...     tree,
... )
{
    'identity': {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)},
    'double': (
        {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)},
        {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)}
    )
}
>>> tree_transpose_map(  
...     lambda x: {'identity': x, 'double': (x, x)},
...     tree,
...     inner_treespec=tree_structure({'identity': 0, 'double': 0}),
... )
{
    'identity': {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)},
    'double': {'b': ((2, 2), [(3, 3), (4, 4)]), 'a': (1, 1), 'c': ((5, 5), (6, 6))}
}

Parameters:

func (callable) – A function that takes 1 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees.
tree (pytree) – A pytree to be mapped over, with each leaf providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, each of which has the same structure as tree or has tree as a prefix.
inner_treespec (PyTreeSpec, optional) – The treespec object representing the inner structure of the result pytree. If not specified, the inner structure is inferred from the result of the function func on the first leaf. (default: None)
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(U), Tuple[PyTree[U], ...], List[PyTree[U]], Dict[Any, PyTree[U]], Deque[PyTree[U]], CustomTreeNode[PyTree[U]]]

Returns:

A new nested pytree with the same structure as inner_treespec but with the value at each leaf has the same structure as tree. The subtree at each leaf is given by the result of function func(x, *xs) where x is the value at the corresponding leaf in tree and xs is the tuple of values at corresponding nodes in rests.

optree.tree_transpose_map_with_path(func, tree, *rests, inner_treespec=None, is_leaf=None, none_is_leaf=False, namespace='')[source]

Map a multi-input function over pytree args as well as the tree paths to produce a new pytree with transposed structure.

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)}
>>> tree_transpose_map_with_path(  
...     lambda p, x: {'depth': len(p), 'value': x},
...     tree,
... )
{
    'depth': {'b': (2, [3, 3]), 'a': 1, 'c': (2, 2)},
    'value': {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)}
}
>>> tree_transpose_map_with_path(  
...     lambda p, x: {'path': p, 'value': x},
...     tree,
...     inner_treespec=tree_structure({'path': 0, 'value': 0}),
... )
{
    'path': {
        'b': (('b', 0), [('b', 1, 0), ('b', 1, 1)]),
        'a': ('a',),
        'c': (('c', 0), ('c', 1))
    },
    'value': {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)}
}

Parameters:

func (callable) – A function that takes 2 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees with extra paths.
tree (pytree) – A pytree to be mapped over, with each leaf providing the second positional argument and the corresponding path providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, each of which has the same structure as tree or has tree as a prefix.
inner_treespec (PyTreeSpec, optional) – The treespec object representing the inner structure of the result pytree. If not specified, the inner structure is inferred from the result of the function func on the first leaf. (default: None)
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(U), Tuple[PyTree[U], ...], List[PyTree[U]], Dict[Any, PyTree[U]], Deque[PyTree[U]], CustomTreeNode[PyTree[U]]]

Returns:

A new nested pytree with the same structure as inner_treespec but with the value at each leaf has the same structure as tree. The subtree at each leaf is given by the result of function func(p, x, *xs) where (p, x) are the path and value at the corresponding leaf in tree and xs is the tuple of values at corresponding nodes in rests.

optree.tree_transpose_map_with_accessor(func, tree, *rests, inner_treespec=None, is_leaf=None, none_is_leaf=False, namespace='')[source]

Map a multi-input function over pytree args as well as the tree accessors to produce a new pytree with transposed structure.

>>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)}
>>> tree_transpose_map_with_accessor(  
...     lambda a, x: {'depth': len(a), 'code': a.codegen('tree'), 'value': x},
...     tree,
... )
{
    'depth': {
        'b': (2, [3, 3]),
        'a': 1,
        'c': (2, 2)
    },
    'code': {
        'b': ("tree['b'][0]", ["tree['b'][1][0]", "tree['b'][1][1]"]),
        'a': "tree['a']",
        'c': ("tree['c'][0]", "tree['c'][1]")
    },
    'value': {
        'b': (2, [3, 4]),
        'a': 1,
        'c': (5, 6)
    }
}
>>> tree_transpose_map_with_accessor(  
...     lambda a, x: {'path': a.path, 'accessor': a, 'value': x},
...     tree,
...     inner_treespec=tree_structure({'path': 0, 'accessor': 0, 'value': 0}),
... )
{
    'path': {
        'b': (('b', 0), [('b', 1, 0), ('b', 1, 1)]),
        'a': ('a',),
        'c': (('c', 0), ('c', 1))
    },
    'accessor': {
        'b': (
            PyTreeAccessor(*['b'][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
            [
                PyTreeAccessor(*['b'][1][0], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=0, type=<class 'list'>))),
                PyTreeAccessor(*['b'][1][1], (MappingEntry(key='b', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>), SequenceEntry(index=1, type=<class 'list'>)))
            ]
        ),
        'a': PyTreeAccessor(*['a'], (MappingEntry(key='a', type=<class 'dict'>),)),
        'c': (
            PyTreeAccessor(*['c'][0], (MappingEntry(key='c', type=<class 'dict'>), SequenceEntry(index=0, type=<class 'tuple'>))),
            PyTreeAccessor(*['c'][1], (MappingEntry(key='c', type=<class 'dict'>), SequenceEntry(index=1, type=<class 'tuple'>)))
        )
    },
    'value': {'b': (2, [3, 4]), 'a': 1, 'c': (5, 6)}
}

Parameters:

func (callable) – A function that takes 2 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees with extra accessors.
tree (pytree) – A pytree to be mapped over, with each leaf providing the second positional argument and the corresponding path providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, each of which has the same structure as tree or has tree as a prefix.
inner_treespec (PyTreeSpec, optional) – The treespec object representing the inner structure of the result pytree. If not specified, the inner structure is inferred from the result of the function func on the first leaf. (default: None)
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(U), Tuple[PyTree[U], ...], List[PyTree[U]], Dict[Any, PyTree[U]], Deque[PyTree[U]], CustomTreeNode[PyTree[U]]]

Returns:

A new nested pytree with the same structure as inner_treespec but with the value at each leaf has the same structure as tree. The subtree at each leaf is given by the result of function func(a, x, *xs) where (a, x) are the accessor and value at the corresponding leaf in tree and xs is the tuple of values at corresponding nodes in rests.

optree.tree_broadcast_prefix(prefix_tree, full_tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Return a pytree of same structure of full_tree with broadcasted subtrees in prefix_tree.

If a prefix_tree is a prefix of a full_tree, this means the full_tree can be constructed by replacing the leaves of prefix_tree with appropriate subtrees.

This function returns a pytree with the same size as full_tree. The leaves are replicated from prefix_tree. The number of replicas is determined by the corresponding subtree in full_tree.

>>> tree_broadcast_prefix(1, [2, 3, 4])
[1, 1, 1]
>>> tree_broadcast_prefix([1, 2, 3], [4, 5, 6])
[1, 2, 3]
>>> tree_broadcast_prefix([1, 2, 3], [4, 5, 6, 7])
Traceback (most recent call last):
    ...
ValueError: list arity mismatch; expected: 3, got: 4; list: [4, 5, 6, 7].
>>> tree_broadcast_prefix([1, 2, 3], [4, 5, (6, 7)])
[1, 2, (3, 3)]
>>> tree_broadcast_prefix([1, 2, 3], [4, 5, {'a': 6, 'b': 7, 'c': (None, 8)}])
[1, 2, {'a': 3, 'b': 3, 'c': (None, 3)}]
>>> tree_broadcast_prefix([1, 2, 3], [4, 5, {'a': 6, 'b': 7, 'c': (None, 8)}], none_is_leaf=True)
[1, 2, {'a': 3, 'b': 3, 'c': (3, 3)}]

Parameters:

prefix_tree (pytree) – A pytree with the prefix structure of full_tree.
full_tree (pytree) – A pytree with the suffix structure of prefix_tree.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]

Returns:

A pytree of same structure of full_tree with broadcasted subtrees in prefix_tree.

optree.broadcast_prefix(prefix_tree, full_tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Return a list of broadcasted leaves in prefix_tree to match the number of leaves in full_tree.

If a prefix_tree is a prefix of a full_tree, this means the full_tree can be constructed by replacing the leaves of prefix_tree with appropriate subtrees.

This function returns a list of leaves with the same size as full_tree. The leaves are replicated from prefix_tree. The number of replicas is determined by the corresponding subtree in full_tree.

>>> broadcast_prefix(1, [2, 3, 4])
[1, 1, 1]
>>> broadcast_prefix([1, 2, 3], [4, 5, 6])
[1, 2, 3]
>>> broadcast_prefix([1, 2, 3], [4, 5, 6, 7])
Traceback (most recent call last):
    ...
ValueError: list arity mismatch; expected: 3, got: 4; list: [4, 5, 6, 7].
>>> broadcast_prefix([1, 2, 3], [4, 5, (6, 7)])
[1, 2, 3, 3]
>>> broadcast_prefix([1, 2, 3], [4, 5, {'a': 6, 'b': 7, 'c': (None, 8)}])
[1, 2, 3, 3, 3]
>>> broadcast_prefix([1, 2, 3], [4, 5, {'a': 6, 'b': 7, 'c': (None, 8)}], none_is_leaf=True)
[1, 2, 3, 3, 3, 3]

Parameters:

prefix_tree (pytree) – A pytree with the prefix structure of full_tree.
full_tree (pytree) – A pytree with the suffix structure of prefix_tree.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

list[TypeVar(T)]

Returns:

A list of leaves in prefix_tree broadcasted to match the number of leaves in full_tree.

optree.tree_broadcast_common(tree, other_tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Return two pytrees of common suffix structure of tree and other_tree with broadcasted subtrees.

If a suffix_tree is a suffix of a tree, this means the suffix_tree can be constructed by replacing the leaves of tree with appropriate subtrees.

This function returns two pytrees with the same structure. The tree structure is the common suffix structure of tree and other_tree. The leaves are replicated from tree and other_tree. The number of replicas is determined by the corresponding subtree in the suffix structure.

>>> tree_broadcast_common(1, [2, 3, 4])
([1, 1, 1], [2, 3, 4])
>>> tree_broadcast_common([1, 2, 3], [4, 5, 6])
([1, 2, 3], [4, 5, 6])
>>> tree_broadcast_common([1, 2, 3], [4, 5, 6, 7])
Traceback (most recent call last):
    ...
ValueError: list arity mismatch; expected: 3, got: 4.
>>> tree_broadcast_common([1, (2, 3), 4], [5, 6, (7, 8)])
([1, (2, 3), (4, 4)], [5, (6, 6), (7, 8)])
>>> tree_broadcast_common([1, {'a': (2, 3)}, 4], [5, 6, {'a': 7, 'b': 8, 'c': (None, 9)}])
([1, {'a': (2, 3)}, {'a': 4, 'b': 4, 'c': (None, 4)}],
 [5, {'a': (6, 6)}, {'a': 7, 'b': 8, 'c': (None, 9)}])
>>> tree_broadcast_common([1, {'a': (2, 3)}, 4], [5, 6, {'a': 7, 'b': 8, 'c': (None, 9)}], none_is_leaf=True)
([1, {'a': (2, 3)}, {'a': 4, 'b': 4, 'c': (4, 4)}],
 [5, {'a': (6, 6)}, {'a': 7, 'b': 8, 'c': (None, 9)}])
>>> tree_broadcast_common([1, None], [None, 2])
([None, None], [None, None])
>>> tree_broadcast_common([1, None], [None, 2], none_is_leaf=True)
([1, None], [None, 2])

Parameters:

tree (pytree) – A pytree has a common suffix structure of other_tree.
other_tree (pytree) – A pytree has a common suffix structure of tree.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

tuple[Union[TypeVar(T), Tuple[Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]], ...], List[Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]], Dict[Any, Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]], Deque[Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]], CustomTreeNode[Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]]], Union[TypeVar(T), Tuple[Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]], ...], List[Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]], Dict[Any, Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]], Deque[Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]], CustomTreeNode[Union[TypeVar(T), Tuple[PyTree[T], ...], List[PyTree[T]], Dict[Any, PyTree[T]], Deque[PyTree[T]], CustomTreeNode[PyTree[T]]]]]]

Returns:

Two pytrees of common suffix structure of tree and other_tree with broadcasted subtrees.

optree.broadcast_common(tree, other_tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Return two lists of leaves in tree and other_tree broadcasted to match the number of leaves in the common suffix structure.

If a suffix_tree is a suffix of a tree, this means the suffix_tree can be constructed by replacing the leaves of tree with appropriate subtrees.

This function returns two pytrees with the same structure. The tree structure is the common suffix structure of tree and other_tree. The leaves are replicated from tree and other_tree. The number of replicas is determined by the corresponding subtree in the suffix structure.

>>> broadcast_common(1, [2, 3, 4])
([1, 1, 1], [2, 3, 4])
>>> broadcast_common([1, 2, 3], [4, 5, 6])
([1, 2, 3], [4, 5, 6])
>>> broadcast_common([1, 2, 3], [4, 5, 6, 7])
Traceback (most recent call last):
    ...
ValueError: list arity mismatch; expected: 3, got: 4.
>>> broadcast_common([1, (2, 3), 4], [5, 6, (7, 8)])
([1, 2, 3, 4, 4], [5, 6, 6, 7, 8])
>>> broadcast_common([1, {'a': (2, 3)}, 4], [5, 6, {'a': 7, 'b': 8, 'c': (None, 9)}])
([1, 2, 3, 4, 4, 4], [5, 6, 6, 7, 8, 9])
>>> broadcast_common([1, {'a': (2, 3)}, 4], [5, 6, {'a': 7, 'b': 8, 'c': (None, 9)}], none_is_leaf=True)
([1, 2, 3, 4, 4, 4, 4], [5, 6, 6, 7, 8, None, 9])
>>> broadcast_common([1, None], [None, 2])
([], [])
>>> broadcast_common([1, None], [None, 2], none_is_leaf=True)
([1, None], [None, 2])

Parameters:

tree (pytree) – A pytree has a common suffix structure of other_tree.
other_tree (pytree) – A pytree has a common suffix structure of tree.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

tuple[list[TypeVar(T)], list[TypeVar(T)]]

Returns:

Two lists of leaves in tree and other_tree broadcasted to match the number of leaves in the common suffix structure.

optree.tree_broadcast_map(func, tree, *rests, is_leaf=None, none_is_leaf=False, namespace='')[source]

Map a multi-input function over pytree args to produce a new pytree.

If only one input is provided, this function is the same as tree_map():

>>> tree_broadcast_map(lambda x: x + 1, {'x': 7, 'y': (42, 64)})
{'x': 8, 'y': (43, 65)}
>>> tree_broadcast_map(lambda x: x + 1, {'x': 7, 'y': (42, 64), 'z': None})
{'x': 8, 'y': (43, 65), 'z': None}
>>> tree_broadcast_map(lambda x: x is None, {'x': 7, 'y': (42, 64), 'z': None})
{'x': False, 'y': (False, False), 'z': None}
>>> tree_broadcast_map(lambda x: x is None, {'x': 7, 'y': (42, 64), 'z': None}, none_is_leaf=True)
{'x': False, 'y': (False, False), 'z': True}

If multiple inputs are given, all input trees will be broadcasted to the common suffix structure of all inputs:

>>> tree_broadcast_map(lambda x, y: x * y, [5, 6, (3, 4)], [{'a': 7, 'b': 9}, [1, 2], 8])
[{'a': 35, 'b': 45}, [6, 12], (24, 32)]

Parameters:

func (callable) – A function that takes 1 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees.
tree (pytree) – A pytree to be mapped over, with each leaf providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, they should have a common suffix structure with each other and with tree.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(U), Tuple[PyTree[U], ...], List[PyTree[U]], Dict[Any, PyTree[U]], Deque[PyTree[U]], CustomTreeNode[PyTree[U]]]

Returns:

A new pytree with the structure as the common suffix structure of tree and rests but with the value at each leaf given by func(x, *xs) where x is the value at the corresponding leaf (may be broadcasted) in tree and xs is the tuple of values at corresponding leaves (may be broadcasted) in rests.

optree.tree_broadcast_map_with_path(func, tree, *rests, is_leaf=None, none_is_leaf=False, namespace='')[source]

Map a multi-input function over pytree args as well as the tree paths to produce a new pytree.

If only one input is provided, this function is the same as tree_map():

>>> tree_broadcast_map_with_path(lambda p, x: (len(p), x), {'x': 7, 'y': (42, 64)})
{'x': (1, 7), 'y': ((2, 42), (2, 64))}
>>> tree_broadcast_map_with_path(lambda p, x: x + len(p), {'x': 7, 'y': (42, 64), 'z': None})
{'x': 8, 'y': (44, 66), 'z': None}
>>> tree_broadcast_map_with_path(lambda p, x: p, {'x': 7, 'y': (42, 64), 'z': {1.5: None}})
{'x': ('x',), 'y': (('y', 0), ('y', 1)), 'z': {1.5: None}}
>>> tree_broadcast_map_with_path(lambda p, x: p, {'x': 7, 'y': (42, 64), 'z': {1.5: None}}, none_is_leaf=True)
{'x': ('x',), 'y': (('y', 0), ('y', 1)), 'z': {1.5: ('z', 1.5)}}

If multiple inputs are given, all input trees will be broadcasted to the common suffix structure of all inputs:

>>> tree_broadcast_map_with_path(  
...     lambda p, x, y: (p, x * y),
...     [5, 6, (3, 4)],
...     [{'a': 7, 'b': 9}, [1, 2], 8],
... )
[
    {'a': ((0, 'a'), 35), 'b': ((0, 'b'), 45)},
    [((1, 0), 6), ((1, 1), 12)],
    (((2, 0), 24), ((2, 1), 32))
]

Parameters:

func (callable) – A function that takes 2 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees with extra paths.
tree (pytree) – A pytree to be mapped over, with each leaf providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, they should have a common suffix structure with each other and with tree.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(U), Tuple[PyTree[U], ...], List[PyTree[U]], Dict[Any, PyTree[U]], Deque[PyTree[U]], CustomTreeNode[PyTree[U]]]

Returns:

A new pytree with the structure as the common suffix structure of tree and rests but with the value at each leaf given by func(p, x, *xs) where (p, x) are the path and value at the corresponding leaf (may be broadcasted) in and xs is the tuple of values at corresponding leaves (may be broadcasted) in rests.

optree.tree_broadcast_map_with_accessor(func, tree, *rests, is_leaf=None, none_is_leaf=False, namespace='')[source]

Map a multi-input function over pytree args as well as the tree accessors to produce a new pytree.

If only one input is provided, this function is the same as tree_map():

>>> tree_broadcast_map_with_accessor(lambda a, x: (len(a), x), {'x': 7, 'y': (42, 64)})
{'x': (1, 7), 'y': ((2, 42), (2, 64))}
>>> tree_broadcast_map_with_accessor(lambda a, x: x + len(a), {'x': 7, 'y': (42, 64), 'z': None})
{'x': 8, 'y': (44, 66), 'z': None}
>>> tree_broadcast_map_with_accessor(  
...     lambda a, x: a.codegen('tree'),
...     {'x': 7, 'y': (42, 64), 'z': {1.5: None}},
... )
{
    'x': "tree['x']",
    'y': ("tree['y'][0]", "tree['y'][1]"),
    'z': {1.5: None}
}
>>> tree_broadcast_map_with_accessor(  
...     lambda a, x: a.codegen('tree'),
...     {'x': 7, 'y': (42, 64), 'z': {1.5: None}},
...     none_is_leaf=True,
... )
{
    'x': "tree['x']",
    'y': ("tree['y'][0]", "tree['y'][1]"),
    'z': {1.5: "tree['z'][1.5]"}
}

If multiple inputs are given, all input trees will be broadcasted to the common suffix structure of all inputs:

>>> tree_broadcast_map_with_accessor(  
...     lambda a, x, y: f'{a.codegen("tree")} = {x * y}',
...     [5, 6, (3, 4)],
...     [{'a': 7, 'b': 9}, [1, 2], 8],
... )
[
    {'a': "tree[0]['a'] = 35", 'b': "tree[0]['b'] = 45"},
    ['tree[1][0] = 6', 'tree[1][1] = 12'],
    ('tree[2][0] = 24', 'tree[2][1] = 32')
]

Parameters:

func (callable) – A function that takes 2 + len(rests) arguments, to be applied at the corresponding leaves of the pytrees with extra accessors.
tree (pytree) – A pytree to be mapped over, with each leaf providing the first positional argument to function func.
rests (tuple of pytree) – A tuple of pytrees, they should have a common suffix structure with each other and with tree.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Union[TypeVar(U), Tuple[PyTree[U], ...], List[PyTree[U]], Dict[Any, PyTree[U]], Deque[PyTree[U]], CustomTreeNode[PyTree[U]]]

Returns:

A new pytree with the structure as the common suffix structure of tree and rests but with the value at each leaf given by func(a, x, *xs) where (a, x) are the accessor and value at the corresponding leaf (may be broadcasted) in and xs is the tuple of values at corresponding leaves (may be broadcasted) in rests.

optree.tree_flatten_one_level(tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Flatten the pytree one level, returning a 4-tuple of children, auxiliary data, path entries, and an unflatten function.

>>> children, metadata, entries, unflatten_func = tree_flatten_one_level({'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5})
>>> children, metadata, entries
([1, (2, [3, 4]), None, 5], ['a', 'b', 'c', 'd'], ('a', 'b', 'c', 'd'))
>>> unflatten_func(metadata, children)
{'a': 1, 'b': (2, [3, 4]), 'c': None, 'd': 5}
>>> children, metadata, entries, unflatten_func = tree_flatten_one_level([{'a': 1, 'b': (2, 3)}, (4, 5)])
>>> children, metadata, entries
([{'a': 1, 'b': (2, 3)}, (4, 5)], None, (0, 1))
>>> unflatten_func(metadata, children)
[{'a': 1, 'b': (2, 3)}, (4, 5)]

Parameters:

tree (pytree) – A pytree to be traversed.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

Returns:

A 4-tuple (children, metadata, entries, unflatten_func). The first element is a list of one-level children of the pytree node. The second element is the auxiliary data used to reconstruct the pytree node. The third element is a tuple of path entries to the children. The fourth element is a function that can be used to unflatten the auxiliary data and children back to the pytree node.

optree.prefix_errors(prefix_tree, full_tree, is_leaf=None, *, none_is_leaf=False, namespace='')[source]

Return a list of errors that would be raised by broadcast_prefix().

Return type:: list[Callable[[str], ValueError]]

Tree Reduce Functions

`tree_reduce`(func, tree[, initial, is_leaf, ...])	Traversal through a pytree and reduce the leaves in left-to-right depth-first order.
`tree_sum`(tree[, start, is_leaf, ...])	Sum `start` and leaf values in `tree` in left-to-right depth-first order and return the total.
`tree_max`(tree, *[, default, key, is_leaf, ...])	Return the maximum leaf value in `tree`.
`tree_min`(tree, *[, default, key, is_leaf, ...])	Return the minimum leaf value in `tree`.
`tree_all`(tree, *[, is_leaf, none_is_leaf, ...])	Test whether all leaves in `tree` are true (or if `tree` is empty).
`tree_any`(tree, *[, is_leaf, none_is_leaf, ...])	Test whether all leaves in `tree` are true (or `False` if `tree` is empty).

optree.tree_reduce(func, tree, initial=<MISSING>, *, is_leaf=None, none_is_leaf=False, namespace='')[source]

Traversal through a pytree and reduce the leaves in left-to-right depth-first order.

See also tree_leaves() and tree_sum().

>>> tree_reduce(lambda x, y: x + y, {'x': 1, 'y': (2, 3)})
6
>>> tree_reduce(lambda x, y: x + y, {'x': 1, 'y': (2, None), 'z': 3})  # `None` is a non-leaf node with arity 0 by default
6
>>> tree_reduce(lambda x, y: x and y, {'x': 1, 'y': (2, None), 'z': 3})
3
>>> tree_reduce(lambda x, y: x and y, {'x': 1, 'y': (2, None), 'z': 3}, none_is_leaf=True)
None

Parameters:

func (callable) – A function that takes two arguments and returns a value of the same type.
tree (pytree) – A pytree to be traversed.
initial (object, optional) – An initial value to be used for the reduction. If not provided, the first leaf value is used as the initial value.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

TypeVar(T)

Returns:

The result of reducing the leaves of the pytree using func.

optree.tree_sum(tree, start=0, *, is_leaf=None, none_is_leaf=False, namespace='')[source]

Sum start and leaf values in tree in left-to-right depth-first order and return the total.

See also tree_leaves() and tree_reduce().

>>> tree_sum({'x': 1, 'y': (2, 3)})
6
>>> tree_sum({'x': 1, 'y': (2, None), 'z': 3})  # `None` is a non-leaf node with arity 0 by default
6
>>> tree_sum({'x': 1, 'y': (2, None), 'z': 3}, none_is_leaf=True)
Traceback (most recent call last):
    ...
TypeError: unsupported operand type(s) for +: 'int' and 'NoneType'
>>> tree_sum({'x': 'a', 'y': ('b', None), 'z': 'c'}, start='')
'abc'
>>> tree_sum({'x': [1], 'y': ([2], [None]), 'z': [3]}, start=[], is_leaf=lambda x: isinstance(x, list))
[1, 2, None, 3]

Parameters:

tree (pytree) – A pytree to be traversed.
start (object, optional) – An initial value to be used for the sum. (default: 0)
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

TypeVar(T)

Returns:

The total sum of start and leaf values in tree.

optree.tree_max(tree, *, default=<MISSING>, key=None, is_leaf=None, none_is_leaf=False, namespace='')[source]

Return the maximum leaf value in tree.

See also tree_leaves() and tree_min().

>>> tree_max({})
Traceback (most recent call last):
    ...
ValueError: max() iterable argument is empty
>>> tree_max({}, default=0)
0
>>> tree_max({'x': 0, 'y': (2, 1)})
2
>>> tree_max({'x': 0, 'y': (2, 1)}, key=lambda x: -x)
0
>>> tree_max({'a': None})  # `None` is a non-leaf node with arity 0 by default
Traceback (most recent call last):
    ...
ValueError: max() iterable argument is empty
>>> tree_max({'a': None}, default=0)  # `None` is a non-leaf node with arity 0 by default
0
>>> tree_max({'a': None}, none_is_leaf=True)
None
>>> tree_max(None)  # `None` is a non-leaf node with arity 0 by default
Traceback (most recent call last):
    ...
ValueError: max() iterable argument is empty
>>> tree_max(None, default=0)
0
>>> tree_max(None, none_is_leaf=True)
None

Parameters:

tree (pytree) – A pytree to be traversed.
default (object, optional) – The default value to return if tree is empty. If the tree is empty and default is not specified, raise a ValueError.
key (callable or None, optional) – An one argument ordering function like that used for list.sort().
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

TypeVar(T)

Returns:

The maximum leaf value in tree.

optree.tree_min(tree, *, default=<MISSING>, key=None, is_leaf=None, none_is_leaf=False, namespace='')[source]

Return the minimum leaf value in tree.

See also tree_leaves() and tree_max().

>>> tree_min({})
Traceback (most recent call last):
    ...
ValueError: min() iterable argument is empty
>>> tree_min({}, default=0)
0
>>> tree_min({'x': 0, 'y': (2, 1)})
0
>>> tree_min({'x': 0, 'y': (2, 1)}, key=lambda x: -x)
2
>>> tree_min({'a': None})  # `None` is a non-leaf node with arity 0 by default
Traceback (most recent call last):
    ...
ValueError: min() iterable argument is empty
>>> tree_min({'a': None}, default=0)  # `None` is a non-leaf node with arity 0 by default
0
>>> tree_min({'a': None}, none_is_leaf=True)
None
>>> tree_min(None)  # `None` is a non-leaf node with arity 0 by default
Traceback (most recent call last):
    ...
ValueError: min() iterable argument is empty
>>> tree_min(None, default=0)
0
>>> tree_min(None, none_is_leaf=True)
None

Parameters:

tree (pytree) – A pytree to be traversed.
default (object, optional) – The default value to return if tree is empty. If the tree is empty and default is not specified, raise a ValueError.
key (callable or None, optional) – An one argument ordering function like that used for list.sort().
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

TypeVar(T)

Returns:

The minimum leaf value in tree.

optree.tree_all(tree, *, is_leaf=None, none_is_leaf=False, namespace='')[source]

Test whether all leaves in tree are true (or if tree is empty).

See also tree_leaves() and tree_any().

>>> tree_all({})
True
>>> tree_all({'x': 1, 'y': (2, 3)})
True
>>> tree_all({'x': 1, 'y': (2, None), 'z': 3})  # `None` is a non-leaf node by default
True
>>> tree_all({'x': 1, 'y': (2, None), 'z': 3}, none_is_leaf=True)
False
>>> tree_all(None)  # `None` is a non-leaf node by default
True
>>> tree_all(None, none_is_leaf=True)
False

Parameters:

tree (pytree) – A pytree to be traversed.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

bool

Returns:

True if all leaves in tree are true, or if tree is empty. Otherwise, False.

optree.tree_any(tree, *, is_leaf=None, none_is_leaf=False, namespace='')[source]

Test whether all leaves in tree are true (or False if tree is empty).

See also tree_leaves() and tree_all().

>>> tree_any({})
False
>>> tree_any({'x': 0, 'y': (2, 0)})
True
>>> tree_any({'a': None})  # `None` is a non-leaf node with arity 0 by default
False
>>> tree_any({'a': None}, none_is_leaf=True)  # `None` is evaluated as false
False
>>> tree_any(None)  # `None` is a non-leaf node with arity 0 by default
False
>>> tree_any(None, none_is_leaf=True)  # `None` is evaluated as false
False

Parameters:

tree (pytree) – A pytree to be traversed.
is_leaf (callable, optional) – An optionally specified function that will be called at each flattening step. It should return a boolean, with True stopping the traversal and the whole subtree being treated as a leaf, and False indicating the flattening should traverse the current object.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

bool

Returns:

True if any leaves in tree are true, otherwise, False. If tree is empty, return False.

PyTreeSpec Functions

`treespec_paths`(treespec)	Return a list of paths to the leaves of a treespec.
`treespec_accessors`(treespec)	Return a list of accessors to the leaves of a treespec.
`treespec_entries`(treespec)	Return a list of one-level entries of a treespec to its children.
`treespec_entry`(treespec, index)	Return the entry of a treespec at the given index.
`treespec_children`(treespec)	Return a list of treespecs for the children of a treespec.
`treespec_child`(treespec, index)	Return the treespec of the child of a treespec at the given index.
`treespec_is_leaf`(treespec[, strict])	Return whether the treespec is a leaf that has no children.
`treespec_is_strict_leaf`(treespec)	Return whether the treespec is a strict leaf.
`treespec_is_prefix`(treespec, other_treespec)	Return whether `treespec` is a prefix of `other_treespec`.
`treespec_is_suffix`(treespec, other_treespec)	Return whether `treespec` is a suffix of `other_treespec`.
`treespec_leaf`(*[, none_is_leaf, namespace])	Make a treespec representing a leaf node.
`treespec_none`(*[, none_is_leaf, namespace])	Make a treespec representing a `None` node.
`treespec_tuple`([iterable, none_is_leaf, ...])	Make a tuple treespec from a list of child treespecs.
`treespec_list`([iterable, none_is_leaf, ...])	Make a list treespec from a list of child treespecs.
`treespec_dict`([mapping, none_is_leaf, namespace])	Make a dict treespec from a dict of child treespecs.
`treespec_namedtuple`(namedtuple, *[, ...])	Make a namedtuple treespec from a namedtuple of child treespecs.
`treespec_ordereddict`([mapping, ...])	Make an OrderedDict treespec from an OrderedDict of child treespecs.
`treespec_defaultdict`([default_factory, ...])	Make a defaultdict treespec from a defaultdict of child treespecs.
`treespec_deque`([iterable, maxlen, ...])	Make a deque treespec from a deque of child treespecs.
`treespec_structseq`(structseq, *[, ...])	Make a PyStructSequence treespec from a PyStructSequence of child treespecs.
`treespec_from_collection`(collection, *[, ...])	Make a treespec from a collection of child treespecs.

optree.treespec_paths(treespec)[source]

Return a list of paths to the leaves of a treespec.

Return type:: list[tuple[Any, ...]]

optree.treespec_accessors(treespec)[source]

Return a list of accessors to the leaves of a treespec.

Return type:: list[PyTreeAccessor]

optree.treespec_entries(treespec)[source]

Return a list of one-level entries of a treespec to its children.

Return type:: list[Any]

optree.treespec_entry(treespec, index)[source]

Return the entry of a treespec at the given index.

Return type:: Any

optree.treespec_children(treespec)[source]

Return a list of treespecs for the children of a treespec.

Return type:: list[PyTreeSpec]

optree.treespec_child(treespec, index)[source]

Return the treespec of the child of a treespec at the given index.

Return type:: PyTreeSpec

optree.treespec_is_leaf(treespec, strict=True)[source]

Return whether the treespec is a leaf that has no children.

This function is equivalent to treespec.is_leaf(strict=strict). If strict=False, it will return True if and only if the treespec represents a strict leaf. If strict=False, it will return True if the treespec represents a strict leaf or None or an empty container (e.g., an empty tuple).

>>> treespec_is_leaf(tree_structure(1))
True
>>> treespec_is_leaf(tree_structure((1, 2)))
False
>>> treespec_is_leaf(tree_structure(None))
False
>>> treespec_is_leaf(tree_structure(None), strict=False)
True
>>> treespec_is_leaf(tree_structure(None, none_is_leaf=False))
False
>>> treespec_is_leaf(tree_structure(None, none_is_leaf=True))
True
>>> treespec_is_leaf(tree_structure(()))
False
>>> treespec_is_leaf(tree_structure(()), strict=False)
True
>>> treespec_is_leaf(tree_structure([]))
False
>>> treespec_is_leaf(tree_structure([]), strict=False)
True

Parameters:

treespec (PyTreeSpec) – A treespec.
strict (bool, optional) – Whether not to treat None or an empty container (e.g., an empty tuple) as a leaf. (default: True)

Return type:

bool

Returns:

True if the treespec represents a leaf that has no children, otherwise, False.

optree.treespec_is_strict_leaf(treespec)[source]

Return whether the treespec is a strict leaf.

This function respects the none_is_leaf setting in the treespec. It is equivalent to treespec.is_leaf(strict=True). It will return True if and only if the treespec represents a strict leaf.

>>> treespec_is_strict_leaf(tree_structure(1))
True
>>> treespec_is_strict_leaf(tree_structure((1, 2)))
False
>>> treespec_is_strict_leaf(tree_structure(None))
False
>>> treespec_is_strict_leaf(tree_structure(None, none_is_leaf=False))
False
>>> treespec_is_strict_leaf(tree_structure(None, none_is_leaf=True))
True
>>> treespec_is_strict_leaf(tree_structure(()))
False
>>> treespec_is_strict_leaf(tree_structure([]))
False

Parameters:: treespec (PyTreeSpec) – A treespec.
Return type:: bool
Returns:: True if the treespec represents a strict leaf, otherwise, False.

optree.treespec_is_prefix(treespec, other_treespec, strict=False)[source]

Return whether treespec is a prefix of other_treespec.

Return type:: bool

optree.treespec_is_suffix(treespec, other_treespec, strict=False)[source]

Return whether treespec is a suffix of other_treespec.

Return type:: bool

optree.treespec_leaf(*, none_is_leaf=False, namespace='')[source]

Make a treespec representing a leaf node.

>>> treespec_leaf()
PyTreeSpec(*)
>>> treespec_leaf(none_is_leaf=True)
PyTreeSpec(*, NoneIsLeaf)
>>> treespec_leaf(none_is_leaf=False) == treespec_leaf(none_is_leaf=True)
False
>>> treespec_leaf() == tree_structure(1)
True
>>> treespec_leaf(none_is_leaf=True) == tree_structure(1, none_is_leaf=True)
True
>>> treespec_leaf(none_is_leaf=True) == tree_structure(None, none_is_leaf=True)
True
>>> treespec_leaf(none_is_leaf=True) == tree_structure(None, none_is_leaf=False)
False
>>> treespec_leaf(none_is_leaf=True) == treespec_none(none_is_leaf=True)
True
>>> treespec_leaf(none_is_leaf=True) == treespec_none(none_is_leaf=False)
False
>>> treespec_leaf(none_is_leaf=False) == treespec_none(none_is_leaf=True)
False
>>> treespec_leaf(none_is_leaf=False) == treespec_none(none_is_leaf=False)
False

Parameters:

none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing a leaf node.

optree.treespec_none(*, none_is_leaf=False, namespace='')[source]

Make a treespec representing a None node.

>>> treespec_none()
PyTreeSpec(None)
>>> treespec_none(none_is_leaf=True)
PyTreeSpec(*, NoneIsLeaf)
>>> treespec_none(none_is_leaf=False) == treespec_none(none_is_leaf=True)
False
>>> treespec_none() == tree_structure(None)
True
>>> treespec_none() == tree_structure(1)
False
>>> treespec_none(none_is_leaf=True) == tree_structure(1, none_is_leaf=True)
True
>>> treespec_none(none_is_leaf=True) == tree_structure(None, none_is_leaf=True)
True
>>> treespec_none(none_is_leaf=True) == tree_structure(None, none_is_leaf=False)
False
>>> treespec_none(none_is_leaf=True) == treespec_leaf(none_is_leaf=True)
True
>>> treespec_none(none_is_leaf=False) == treespec_leaf(none_is_leaf=True)
False
>>> treespec_none(none_is_leaf=True) == treespec_leaf(none_is_leaf=False)
False
>>> treespec_none(none_is_leaf=False) == treespec_leaf(none_is_leaf=False)
False

Parameters:

none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing a None node.

optree.treespec_tuple(iterable=(), *, none_is_leaf=False, namespace='')[source]

Make a tuple treespec from a list of child treespecs.

>>> treespec_tuple([treespec_leaf(), treespec_leaf()])
PyTreeSpec((*, *))
>>> treespec_tuple([treespec_leaf(), treespec_leaf(), treespec_none()])
PyTreeSpec((*, *, None))
>>> treespec_tuple()
PyTreeSpec(())
>>> treespec_tuple([treespec_leaf(), treespec_tuple([treespec_leaf(), treespec_leaf()])])
PyTreeSpec((*, (*, *)))
>>> treespec_tuple([treespec_leaf(), tree_structure({'a': 1, 'b': 2})])
PyTreeSpec((*, {'a': *, 'b': *}))
>>> treespec_tuple([treespec_leaf(), tree_structure({'a': 1, 'b': 2}, none_is_leaf=True)])
Traceback (most recent call last):
    ...
ValueError: Expected treespec(s) with `node_is_leaf=False`.

Parameters:

iterable (iterable of PyTreeSpec, optional) – A iterable of child treespecs. They must have the same node_is_leaf and namespace values.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing a tuple node with the given children.

optree.treespec_list(iterable=(), *, none_is_leaf=False, namespace='')[source]

Make a list treespec from a list of child treespecs.

>>> treespec_list([treespec_leaf(), treespec_leaf()])
PyTreeSpec([*, *])
>>> treespec_list([treespec_leaf(), treespec_leaf(), treespec_none()])
PyTreeSpec([*, *, None])
>>> treespec_list()
PyTreeSpec([])
>>> treespec_list([treespec_leaf(), treespec_tuple([treespec_leaf(), treespec_leaf()])])
PyTreeSpec([*, (*, *)])
>>> treespec_list([treespec_leaf(), tree_structure({'a': 1, 'b': 2})])
PyTreeSpec([*, {'a': *, 'b': *}])
>>> treespec_list([treespec_leaf(), tree_structure({'a': 1, 'b': 2}, none_is_leaf=True)])
Traceback (most recent call last):
    ...
ValueError: Expected treespec(s) with `node_is_leaf=False`.

Parameters:

iterable (iterable of PyTreeSpec, optional) – A iterable of child treespecs. They must have the same node_is_leaf and namespace values.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing a list node with the given children.

optree.treespec_dict(mapping=(), *, none_is_leaf=False, namespace='', **kwargs)[source]

Make a dict treespec from a dict of child treespecs.

>>> treespec_dict({'a': treespec_leaf(), 'b': treespec_leaf()})
PyTreeSpec({'a': *, 'b': *})
>>> treespec_dict([('b', treespec_leaf()), ('c', treespec_leaf()), ('a', treespec_none())])
PyTreeSpec({'a': None, 'b': *, 'c': *})
>>> treespec_dict()
PyTreeSpec({})
>>> treespec_dict(a=treespec_leaf(), b=treespec_tuple([treespec_leaf(), treespec_leaf()]))
PyTreeSpec({'a': *, 'b': (*, *)})
>>> treespec_dict({'a': treespec_leaf(), 'b': tree_structure([1, 2])})
PyTreeSpec({'a': *, 'b': [*, *]})
>>> treespec_dict({'a': treespec_leaf(), 'b': tree_structure([1, 2], none_is_leaf=True)})
Traceback (most recent call last):
    ...
ValueError: Expected treespec(s) with `node_is_leaf=False`.

Parameters:

mapping (mapping of PyTreeSpec, optional) – A mapping of child treespecs. They must have the same node_is_leaf and namespace values.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing a dict node with the given children.

optree.treespec_namedtuple(namedtuple, *, none_is_leaf=False, namespace='')[source]

Make a namedtuple treespec from a namedtuple of child treespecs.

>>> from collections import namedtuple
>>> Point = namedtuple('Point', ['x', 'y'])
>>> treespec_namedtuple(Point(x=treespec_leaf(), y=treespec_leaf()))
PyTreeSpec(Point(x=*, y=*))
>>> treespec_namedtuple(Point(x=treespec_leaf(), y=treespec_tuple([treespec_leaf(), treespec_leaf()])))
PyTreeSpec(Point(x=*, y=(*, *)))
>>> treespec_namedtuple(Point(x=treespec_leaf(), y=tree_structure([1, 2])))
PyTreeSpec(Point(x=*, y=[*, *]))
>>> treespec_namedtuple(Point(x=treespec_leaf(), y=tree_structure([1, 2], none_is_leaf=True)))
Traceback (most recent call last):
    ...
ValueError: Expected treespec(s) with `node_is_leaf=False`.

Parameters:

namedtuple (namedtuple of PyTreeSpec) – A namedtuple of child treespecs. They must have the same node_is_leaf and namespace values.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Returns:

A treespec representing a dict node with the given children.

optree.treespec_ordereddict(mapping=(), *, none_is_leaf=False, namespace='', **kwargs)[source]

Make an OrderedDict treespec from an OrderedDict of child treespecs.

>>> treespec_ordereddict({'a': treespec_leaf(), 'b': treespec_leaf()})
PyTreeSpec(OrderedDict({'a': *, 'b': *}))
>>> treespec_ordereddict([('b', treespec_leaf()), ('c', treespec_leaf()), ('a', treespec_none())])
PyTreeSpec(OrderedDict({'b': *, 'c': *, 'a': None}))
>>> treespec_ordereddict()
PyTreeSpec(OrderedDict())
>>> treespec_ordereddict(a=treespec_leaf(), b=treespec_tuple([treespec_leaf(), treespec_leaf()]))
PyTreeSpec(OrderedDict({'a': *, 'b': (*, *)}))
>>> treespec_ordereddict({'a': treespec_leaf(), 'b': tree_structure([1, 2])})
PyTreeSpec(OrderedDict({'a': *, 'b': [*, *]}))
>>> treespec_ordereddict({'a': treespec_leaf(), 'b': tree_structure([1, 2], none_is_leaf=True)})
Traceback (most recent call last):
    ...
ValueError: Expected treespec(s) with `node_is_leaf=False`.

Parameters:

mapping (mapping of PyTreeSpec, optional) – A mapping of child treespecs. They must have the same node_is_leaf and namespace values.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing an OrderedDict node with the given children.

optree.treespec_defaultdict(default_factory=None, mapping=(), *, none_is_leaf=False, namespace='', **kwargs)[source]

Make a defaultdict treespec from a defaultdict of child treespecs.

>>> treespec_defaultdict(int, {'a': treespec_leaf(), 'b': treespec_leaf()})
PyTreeSpec(defaultdict(<class 'int'>, {'a': *, 'b': *}))
>>> treespec_defaultdict(int, [('b', treespec_leaf()), ('c', treespec_leaf()), ('a', treespec_none())])
PyTreeSpec(defaultdict(<class 'int'>, {'a': None, 'b': *, 'c': *}))
>>> treespec_defaultdict()
PyTreeSpec(defaultdict(None, {}))
>>> treespec_defaultdict(int)
PyTreeSpec(defaultdict(<class 'int'>, {}))
>>> treespec_defaultdict(int, a=treespec_leaf(), b=treespec_tuple([treespec_leaf(), treespec_leaf()]))
PyTreeSpec(defaultdict(<class 'int'>, {'a': *, 'b': (*, *)}))
>>> treespec_defaultdict(int, {'a': treespec_leaf(), 'b': tree_structure([1, 2])})
PyTreeSpec(defaultdict(<class 'int'>, {'a': *, 'b': [*, *]}))
>>> treespec_defaultdict(int, {'a': treespec_leaf(), 'b': tree_structure([1, 2], none_is_leaf=True)})
Traceback (most recent call last):
    ...
ValueError: Expected treespec(s) with `node_is_leaf=False`.

Parameters:

default_factory (callable or None, optional) – A factory function that will be used to create a missing value. (default: None)
mapping (mapping of PyTreeSpec, optional) – A mapping of child treespecs. They must have the same node_is_leaf and namespace values.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing a defaultdict node with the given children.

optree.treespec_deque(iterable=(), maxlen=None, *, none_is_leaf=False, namespace='')[source]

Make a deque treespec from a deque of child treespecs.

>>> treespec_deque([treespec_leaf(), treespec_leaf()])
PyTreeSpec(deque([*, *]))
>>> treespec_deque([treespec_leaf(), treespec_leaf(), treespec_none()], maxlen=5)
PyTreeSpec(deque([*, *, None], maxlen=5))
>>> treespec_deque()
PyTreeSpec(deque([]))
>>> treespec_deque([treespec_leaf(), treespec_tuple([treespec_leaf(), treespec_leaf()])])
PyTreeSpec(deque([*, (*, *)]))
>>> treespec_deque([treespec_leaf(), tree_structure({'a': 1, 'b': 2})], maxlen=5)
PyTreeSpec(deque([*, {'a': *, 'b': *}], maxlen=5))
>>> treespec_deque([treespec_leaf(), tree_structure({'a': 1, 'b': 2}, none_is_leaf=True)], maxlen=5)
Traceback (most recent call last):
    ...
ValueError: Expected treespec(s) with `node_is_leaf=False`.

Parameters:

iterable (iterable of PyTreeSpec, optional) – A iterable of child treespecs. They must have the same node_is_leaf and namespace values.
maxlen (int or None, optional) – The maximum size of a deque or None if unbounded. (default: None)
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing a deque node with the given children.

optree.treespec_structseq(structseq, *, none_is_leaf=False, namespace='')[source]

Make a PyStructSequence treespec from a PyStructSequence of child treespecs.

Parameters:

structseq (PyStructSequence of PyTreeSpec) – A PyStructSequence of child treespecs. They must have the same node_is_leaf and namespace values.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing a PyStructSequence node with the given children.

optree.treespec_from_collection(collection, *, none_is_leaf=False, namespace='')[source]

Make a treespec from a collection of child treespecs.

>>> treespec_from_collection(None)
PyTreeSpec(None)
>>> treespec_from_collection(None, none_is_leaf=True)
PyTreeSpec(*, NoneIsLeaf)
>>> treespec_from_collection(object())
PyTreeSpec(*)
>>> treespec_from_collection([treespec_leaf(), treespec_none()])
PyTreeSpec([*, None])
>>> treespec_from_collection({'a': treespec_leaf(), 'b': treespec_none()})
PyTreeSpec({'a': *, 'b': None})
>>> treespec_from_collection(deque([treespec_leaf(), tree_structure({'a': 1, 'b': 2})], maxlen=5))
PyTreeSpec(deque([*, {'a': *, 'b': *}], maxlen=5))
>>> treespec_from_collection({'a': treespec_leaf(), 'b': (treespec_leaf(), treespec_none())})
Traceback (most recent call last):
    ...
ValueError: Expected a(n) dict of PyTreeSpec(s), got {'a': PyTreeSpec(*), 'b': (PyTreeSpec(*), PyTreeSpec(None))}.
>>> treespec_from_collection([treespec_leaf(), tree_structure({'a': 1, 'b': 2}, none_is_leaf=True)])
Traceback (most recent call last):
    ...
ValueError: Expected treespec(s) with `node_is_leaf=False`.

Parameters:

collection (collection of PyTreeSpec) – A collection of child treespecs. They must have the same node_is_leaf and namespace values.
none_is_leaf (bool, optional) – Whether to treat None as a leaf. If False, None is a non-leaf node with arity 0. Thus None is contained in the treespec rather than in the leaves list and None will be remain in the result pytree. (default: False)
namespace (str, optional) – The registry namespace used for custom pytree node types. (default: '', i.e., the global namespace)

Return type:

PyTreeSpec

Returns:

A treespec representing the same structure of the collection with the given children.