Optimized PyTree Utilities

Warnings

• breaking Python 3.7 support was dropped in optree v0.14.0. Users on Python 3.7 must upgrade their Python version or use optree versions prior to 0.14.0.
  Fix: Upgrade Python to 3.8+ (preferably 3.9+) or pin optree to `<0.14.0`.
• breaking Deprecated key path APIs and `optree.Partial` were removed in optree v0.15.0. Any code relying on these older APIs will break.
  Fix: Update code to use the current API. Refer to optree documentation for alternatives to key path APIs and `functools.partial` for `optree.Partial`.
• gotcha When registering a custom PyTree node type using `optree.register_pytree_node` or `optree.register_pytree_node_class`, a `namespace` argument is explicitly required. This prevents accidental collisions between different libraries registering the same type with different behaviors in the same Python interpreter.
  Fix: Always provide a unique, non-empty `namespace` string (e.g., 'mylibrary.pytrees') when registering custom PyTree nodes. Example: `optree.register_pytree_node(MyClass, flatten_func, unflatten_func, namespace='my.namespace')`.
• gotcha By default, `None` is treated as a non-leaf node with zero children. This means it's part of the tree structure (`treespec`), not the list of leaves. To treat `None` as a leaf node, you must explicitly pass `none_is_leaf=True` to functions like `tree_flatten` or `tree_map`.
  Fix: If `None` should be processed as a leaf, use `none_is_leaf=True` in relevant functions: `tree_map(func, tree, none_is_leaf=True)`.
• gotcha Custom `flatten_func` implementations for `register_pytree_node` must include a proper termination condition to prevent infinite recursion, especially if the children can be of the same type as the current node. This can lead to a `RecursionError`.
  Fix: Carefully design `flatten_func` to ensure that it eventually produces only leaf nodes or built-in non-leaf types, or explicitly handle recursion depth. The library defines `MAX_RECURSION_DEPTH`.
• breaking Module naming conventions changed in v0.16.0, affecting direct imports. Specifically, `optree.accessor` became `optree.accessors`, `optree.integration` became `optree.integrations`, etc.
  Fix: Update import statements. For example, change `from optree import accessor` to `from optree import accessors`.

Install

• pip install optree Install latest version

Imports

• tree_flatten

  from optree import tree_flatten

• tree_unflatten

  from optree import tree_unflatten

• tree_map

  from optree import tree_map

• register_pytree_node

  from optree import register_pytree_node

• pytree

  import optree.pytree as pt

  While `from optree import pytree` works, using `import optree.pytree as pt` is a common alias for convenience, especially since v0.14.1 for its `tree_*` aliases.

Quickstart

This quickstart demonstrates the core `tree_map` function to apply a transformation to all leaves of a PyTree, and also shows how to flatten a PyTree into leaves and a structure (`treespec`) and then unflatten it back, potentially with modified leaves.

from optree import tree_map

def add_one(x):
    return x + 1

tree = {'a': 1, 'b': [2, 3], 'c': {'d': 4}}
mapped_tree = tree_map(add_one, tree)
print(mapped_tree)

from optree import tree_flatten, tree_unflatten, PyTreeSpec

leaves, treespec = tree_flatten(tree)
print(f"Leaves: {leaves}")
print(f"TreeSpec: {treespec}")

reconstructed_tree = tree_unflatten(treespec, [l * 10 for l in leaves])
print(f"Reconstructed tree with modified leaves: {reconstructed_tree}")