optree

What it is and what it does

OpTree is a library for working with PyTrees—arbitrarily nested Python containers and custom objects—by providing efficient flatten, unflatten, and map operations. It treats built-in types like dict, list, tuple, namedtuple, and OrderedDict as non-leaf nodes that can be traversed, while custom objects are leaves unless explicitly registered. The core operations are tree flattening (which extracts leaves in deterministic order and captures structure) and unflattening (which reconstructs the tree from leaves and structure), with tree_map built from these primitives.

The package is written in C++ with Python bindings for performance, supports registration of custom container types, and includes a namespace system (optree.pytree) for creating library-specific PyTree utilities. It is actively maintained, has no known vulnerabilities, and provides prebuilt wheels for modern Python versions across multiple platforms, making installation straightforward in most cases.

Use it for:

• Flatten nested model parameters or optimizer states in machine learning frameworks for batch processing.
• Traverse and transform deeply nested configuration dictionaries or JSON-like structures uniformly.
• Register custom container types (e.g., dataclasses) to participate in tree operations without manual recursion.
• Build functional programming patterns that operate on entire tree structures with a single map call.
• Implement library-specific PyTree utilities via the reexport API for downstream packages.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

OpTree provides optimized utilities for flattening, unflattening, and mapping over nested Python container structures (PyTrees) like dicts, lists, tuples, and custom types.

Yes. OpTree is worth installing if you work with nested Python structures in machine learning, scientific computing, or functional programming contexts. It is actively maintained, has no security issues, and provides significant convenience and performance over manual tree traversal. The permissive Apache-2.0 license poses no restriction. Medium install friction is acceptable given the availability of prebuilt wheels for common platforms.

Install

optree on PyPI

pip install optree

uv add optree

poetry add optree

Installing optree

Before you install

Medium install friction due to C++ extension compilation from source; however, prebuilt wheels are available for Python 3.9–3.15 across macOS, Linux, and Windows. Active maintenance with recent releases and no known vulnerabilities.

License in practice

Apache-2.0 is permissive, allowing commercial and private use with minimal restrictions; attribution and license notice are required.

Quickstart

pip install optree

import optree

tree = {'b': (2, [3, 4]), 'a': 1}
leaves, spec = optree.tree_flatten(tree)
result = optree.tree_map(lambda x: x**2, tree)

Requires Python 3.9+; C++ compiler and cmake needed only if building from source rather than using prebuilt wheels.

Verify before relying

• Performance characteristics compared to other PyTree implementations (e.g., JAX's tree_util).
• Memory overhead of PyTreeSpec objects for very large or deeply nested structures.
• Compatibility guarantees for custom type registration across OpTree versions.

Package facts

Evidence: optree-0.19.1-cp310-cp310-macosx_10_9_x86_64.whl; optree-0.19.1-cp310-cp310-macosx_11_0_arm64.whl; optree-0.19.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; optree-0.19.1-cp310-cp310-manylinux_2_26_i686.manylinux_2_28_i686.whl; optree-0.19.1-cp310-cp310-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl; optree-0.19.1-cp310-cp310-manylinux_2_26_s390x.manylinux_2_28_s390x.whl; optree-0.19.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; optree-0.19.1-cp310-cp310-manylinux_2_39_riscv64.whl; optree-0.19.1-cp310-cp310-win32.whl; optree-0.19.1-cp310-cp310-win_amd64.whl; optree-0.19.1-cp311-cp311-macosx_10_9_x86_64.whl; optree-0.19.1-cp311-cp311-macosx_11_0_arm64.whl; optree-0.19.1-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; optree-0.19.1-cp311-cp311-manylinux_2_26_i686.manylinux_2_28_i686.whl; optree-0.19.1-cp311-cp311-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl; optree-0.19.1-cp311-cp311-manylinux_2_26_s390x.manylinux_2_28_s390x.whl; optree-0.19.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; optree-0.19.1-cp311-cp311-manylinux_2_39_riscv64.whl; optree-0.19.1-cp311-cp311-win32.whl; optree-0.19.1-cp311-cp311-win_amd64.whl

Keywords: PyTree, Tree Manipulation, Tree Traversal, Functional Programming

Tags

More Utilities packages