optree_pytree_flatten_unflatten_and_tree_map.py

1import optree
2
3# A nested data structure (pytree)
4tree = {'b': [1, 2], 'a': 3, 'c': (4, 5)}
5
6# Flatten the tree: returns a list of leaves and the tree structure (treespec)
7leaves, treespec = optree.tree_flatten(tree)
8# leaves: [3, 1, 2, 4, 5]
9# treespec: PyTreeSpec({'a': *, 'b': [*, *], 'c': (*, *)})
10
11# Reconstruct the tree from leaves and treespec
12reconstructed_tree = optree.tree_unflatten(treespec, leaves)
13assert reconstructed_tree == tree
14
15# Apply a function to every leaf of the tree
16doubled_tree = optree.tree_map(lambda x: x * 2, tree)
17# doubled_tree: {'b': [2, 4], 'a': 6, 'c': (8, 10)}
18
19# Count the number of leaves in the tree
20num_leaves = optree.tree_leaves_count(tree)
21# num_leaves: 5