OpTree

OpTree is a pytree manipulation library from the MetaOPT team that reimplements JAX-style pytree utilities with a C++ core for substantially faster flattening, mapping, and unflattening of arbitrarily nested Python containers (tuples, lists, dicts, namedtuples, dataclasses, and custom registered types). It underpins reproducible, structure-preserving transforms used throughout differentiable programming, meta-learning, and deep-learning pipelines that need to walk or restructure nested parameter and gradient trees.

Beyond the core tree_flatten/tree_map/tree_unflatten primitives, OpTree ships dataclass and functools integration modules, first-class support for registering custom container types with path/accessor tracking, and dedicated interop helpers for NumPy, PyTorch, JAX, and attrs, so libraries can adopt one consistent pytree substrate instead of hand-rolling recursive tree walkers.

What You Get

• C++-accelerated tree_flatten/tree_unflatten/tree_map with familiar JAX-style pytree semantics
• A global and namespaced type registry for registering custom container types as tree nodes
• Path- and accessor-aware traversal (tree_flatten_with_path, tree_flatten_with_accessor, PyTreeAccessor) for introspecting exactly where each leaf sits
• Built-in dataclasses and functools integration modules plus interop helpers for NumPy, PyTorch, JAX, and attrs
• PyTreeSpec objects that are picklable, hashable, and comparable, letting structure be serialized or compared independent of leaf values

Common Use Cases

• Flattening and remapping nested model parameters/gradients in deep-learning training loops
• Implementing custom differentiable-programming libraries that need JAX-style pytree semantics without depending on JAX
• Registering project-specific container types (dataclasses, custom namedtuples) so they flatten/unflatten transparently
• Comparing or serializing nested data structures by their PyTreeSpec independent of leaf values
• Building path-aware tooling (e.g. diffing, logging) over deeply nested config or state objects

Under The Hood

Architecture The entry point optree/__init__.py aggregates submodules (ops, accessors, dataclasses, functools, integrations, pytree, treespec, typing, registry, utils). Core tree operations are implemented in C++ (src/optree.cpp, src/registry.cpp, src/treespec/*.cpp) and exposed through a compiled _C extension, while the Python layer (ops.py, nearly 4,000 lines) provides the full public API, argument validation, and higher-level compositions (tree_map, tree_reduce, broadcast_prefix, tree_transpose) built on lower-level primitives such as one-level flattening, unflattening, and PyTreeSpec construction/comparison/hashing/serialization. registry.py plus src/registry.cpp maintain a global or namespaced type registry mapping container types to flatten/unflatten functions that the core recursion walks, and treespec.py wraps the C++ PyTreeSpec object so tree structure can be handled independently of leaf values. Separation of concerns is clean: Python owns API ergonomics, type stubs, and dataclass/functools integration, while C++ owns performance-critical recursive traversal and structure hashing/serialization; because every tree_* function is composed from flatten+map+unflatten via the spec, the registry/PyTreeSpec abstraction is the one piece nearly everything else depends on.

Tech Stack Python 3.10+ with a C++ core built via pybind11 and CMake, using setuptools as the build backend and distributed as a compiled wheel bundling the _C extension. Optional integration dependencies for NumPy, PyTorch, JAX, and attrs turn OpTree into a shared pytree substrate across ML frameworks. Documentation is built with Sphinx plus sphinx-autoapi/autodoc-typehints and hosted on Read the Docs. Dev tooling includes pre-commit, ruff, mypy, pylint, cpplint, and doc8 for combined Python/C++ linting, with pytest, pytest-cov, and covdefaults for testing, all wired into GitHub Actions build/lint/test workflows.

Code Quality The test suite mirrors the module layout with one file per area (tree ops, treespec, registry, accessors, dataclasses, functools, typing, utils, prefix errors) plus dedicated concurrent and integrations subdirectories and a shared fixtures module, using pytest with parametrized cases and Codecov coverage reporting. A dedicated CI job runs the suite against a Python debug build, an unusually rigorous check for a native-extension package. Type safety is covered by bundled stub files and a py.typed marker, with mypy included in the lint dependency group. Error handling in the ops layer favors explicit, diagnostic exceptions (including difflib-based structure-mismatch messages) over silent failures, and naming follows a consistent tree_*/treespec_* verb-object convention; the C++ side carries its own clang-format/clang-tidy configuration for consistent native code style.

API Design The public API exposes a small, consistent verb-object naming scheme (tree_flatten, tree_map, tree_unflatten, treespec_*) that mirrors JAX’s pytree API while running on a faster C++ core, letting downstream libraries adopt tree utilities with minimal boilerplate via the integration modules for NumPy, PyTorch, JAX, and attrs. A pytree.reexport(namespace=...) helper lets a library vendor a namespaced copy of the API without repeating the namespace argument on every call, and both path-aware (tree_flatten_with_path) and accessor-aware (tree_flatten_with_accessor) traversal variants give callers fine-grained introspection without hand-rolled recursion. Getting started requires no configuration beyond pip install optree, and the extensive README with runnable examples doubles as onboarding documentation.