optree_pytree_flatten_unflatten_and_tree_map_quickstart.py

1import optree
2
3# A complex nested structure
4tree = {'b': [1, 2], 'a': (3, 4), 'c': None}
5
6# Flatten the tree
7leaves, treespec = optree.tree_flatten(tree)
8# leaves = [3, 4, 1, 2, None], treespec = PyTreeSpec({'a': (None, None), 'b': [None, None], 'c': None})
9
10# Unflatten the leaves back to the tree
11optree.tree_unflatten(treespec, leaves)
12# {'b': [1, 2], 'a': (3, 4), 'c': None}
13
14# Apply a function to each leaf
15optree.tree_map(lambda x: x + 1 if x is not None else x, tree)
16# {'b': [2, 3], 'a': (4, 5), 'c': None}