optree

library0.20.0pypypi✓ verified 25d ago

OpTree is an optimized Python library for working with PyTrees, which are arbitrarily nested Python containers. It provides efficient utilities for flattening, unflattening, and mapping functions over tree structures. The current version is 0.19.0, and the library maintains an active development cycle with frequent releases.

data serialization ai-ml

pip install optree

github ↗package ↗docs ↗

INSTALL

IMPORT

SIG · OPTREE

optree

datapythonv0.20.0

Install

1.8s avg

Import

90ms

Disk

18MB

Pass rate

5/ 10

Env Coverage5 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v0.20.0 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

py 3.10–3.95 runs

build_error

glibc

py 3.10–3.95 runs

installs and imports cleanly · install 1.8s · import 0.090s · 20MB

18MB installed

● package 18MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

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

✗ from optree import pytree

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.

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}")

Debug

Known issues

breakingPython 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`.

affects: >=0.14.0

breakingDeprecated 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`.

affects: >=0.15.0

gotchaWhen 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')`.

affects: All versions

gotchaBy 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)`.

affects: All versions

gotchaCustom `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`.

affects: All versions

breakingModule 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`.

affects: >=0.16.0

gotchaBuilding `optree` from source requires C++ build tools (like `g++` or `clang++`) and `cmake` in the environment. This issue commonly arises when installing on minimal distributions (e.g., Alpine Linux) or when pre-built wheels are not available for the specific Python version/platform, forcing a source build.

fix

Ensure that C++ build tools and `cmake` are installed in the environment where `optree` is being built. For example, on Alpine Linux, this typically involves `apk add build-base g++ cmake`. On Debian/Ubuntu, use `apt-get install build-essential g++ cmake`.

affects: All versions

Errors

Common errors & fixes

ImportError: cannot import name 'NamedTuple' from 'typing_extensions'

This error typically occurs when the `typing_extensions` package installed in your environment is an outdated version that does not provide the `NamedTuple` type, which `optree` or its dependencies (like Keras) require.

fix

Upgrade `typing_extensions` to a recent version: `pip install --upgrade typing-extensions`

ModuleNotFoundError: No module named 'optree'

The `optree` library is not installed in your current Python environment.

fix

Install `optree` using pip: `pip install optree`

AttributeError: module 'optree' has no attribute 'dict_insertion_ordered'

This error indicates that the installed `optree` version is too old and lacks features like `dict_insertion_ordered`, which newer versions of dependent libraries (e.g., PyTorch) might expect.

fix

Upgrade `optree` to the latest version: `pip install --upgrade optree`

RecursionError: Maximum recursion depth exceeded during flattening the tree.

This usually happens when a custom PyTree node's `__tree_flatten__` (or `tree_flatten`) method, or the `is_leaf` predicate, is implemented in a way that leads to infinite recursion during tree traversal, without a proper termination condition for subtrees.

fix

Review the `__tree_flatten__` (or `tree_flatten`) implementation for custom PyTree nodes or the `is_leaf` function to ensure that children are correctly identified and that the recursion terminates for leaf nodes.

TypeError: '<' not supported between instances of 'int' and 'str'

This error often occurs when `optree` (which sorts dictionary keys by default for deterministic flattening) encounters a dictionary with heterogeneous keys (e.g., a mix of integers and strings) that cannot be directly compared.

fix

Ensure dictionary keys are of a comparable type or use `collections.OrderedDict` if insertion order is important and keys are heterogeneous. Alternatively, provide a custom `is_leaf` function if specific keys should be treated as leaves without sorting.

Upgrade

Version history

0.20.0latest on PyPI · released Aug 20, 2026

Audit

Dependencies

typing-extensionsrequiredRequired for type hinting support, especially for Python versions older than 3.12.

Agent activity

19 hits · last 30 days

node

Resources

githubhttps://github.com/metaopt/optree ↗

packagehttps://pypi.org/project/optree/ ↗

docshttps://optree.readthedocs.io ↗