Tree

Tree (dm-tree) is a lightweight Python library developed by Google DeepMind for manipulating nested data structures (also called pytrees). It generalizes Python’s built-in map function to operate over arbitrarily nested collections — including lists, tuples, dicts, and custom container types — while preserving their structure. This makes it particularly useful in machine learning pipelines and JAX-based workflows, where complex parameter trees or hierarchical state representations are common. The library provides efficient operations such as flatten, unflatten, and map_structure, enabling users to apply functions to all leaves of a nested structure seamlessly. Backed by a high-performance C++ core, tree is optimized for large-scale, performance-critical applications.