optree_pytree_flatten_unflatten_map_reduce_quickstart.py

1import optree
2
3# A nested data structure (pytree)
4tree = {'a': [1, 2], 'b': 3}
5
6# Flatten the tree
7leaves, treespec = optree.tree_flatten(tree)
8# leaves: [1, 2, 3]
9# treespec: PyTreeSpec({'a': [*, *], 'b': *})
10
11# Unflatten back to the original structure
12original_tree = optree.tree_unflatten(treespec, leaves)
13assert original_tree == tree
14
15# Map a function over the leaves
16doubled_tree = optree.tree_map(lambda x: x * 2, tree)
17# doubled_tree: {'a': [2, 4], 'b': 6}
18
19# Use tree_reduce to sum all leaves
20total_sum = optree.tree_reduce(lambda x, y: x + y, tree)
21# total_sum: 6