optree_pytree_flatten_unflatten_and_map_operations.py

1import optree
2
3# A nested structure containing lists, dictionaries, and tuples
4tree = {
5    'b': (1, 2),
6    'a': [3, 4],
7    'c': {'d': 5, 'e': (6, 7)},
8}
9
10# Flatten the tree: returns a list of leaves and a PyTreeSpec object
11leaves, treespec = optree.tree_flatten(tree)
12# leaves: [3, 4, 1, 2, 5, 6, 7]
13# treespec: PyTreeSpec({'a': [*, *], 'b': (*, *), 'c': {'d': *, 'e': (*, *)}})
14
15# Reconstruct the tree from leaves and treespec
16reconstructed_tree = optree.tree_unflatten(treespec, leaves)
17assert reconstructed_tree == tree
18
19# Apply a function to every leaf (e.g., add 1 to each number)
20new_tree = optree.tree_map(lambda x: x + 1, tree)
21# new_tree: {'b': (2, 3), 'a': [4, 5], 'c': {'d': 6, 'e': (7, 8)}}