For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python module

max.experimental.tree_utils

Provides utilities for taking nested values apart and putting them back.

Import the module as a namespace. Interior nodes are list, tuple, dict, and any class declaring the protocol; everything else is a leaf. __tree_flatten__ returns (children, meta) and __tree_unflatten__ rebuilds.

from max.experimental import tree_utils as tree

class Linear:
    def __init__(self, weight, eps=1e-5):
        self.weight, self.eps = weight, eps

    def __tree_flatten__(self):
        return {"weight": self.weight}, self.eps

    @classmethod
    def __tree_unflatten__(cls, eps, children):
        return cls(children["weight"], eps)

model = [Linear("w0"), Linear("w1")]

tree.paths(model, leaf=str)               # {"0.weight": "w0", "1.weight": "w1"}
tree.map(str.upper, model, leaf=str)      # a fresh model, weights mapped
tree.update(model, {"0.weight": "new"}, leaf=str)   # written in place
flat, treedef = tree.flatten(model, leaf=str)
tree.unflatten(treedef, flat)             # a rebuilt model, eps intact

Declare __tree_empty__(meta) instead of __tree_unflatten__ when the node must exist before its children, and optionally __tree_setattr__(key, value). Every walk takes leaf, saying where it stops, and shared, saying whether a value reachable by two paths is one object or two.

Flatten and rebuild

flatten	Splits tree into its leaves and the structure around them.
unflatten	Rebuilds a tree from a structure and its leaves.

Read

leaves	Returns tree's leaves, left to right, dropping the structure.
nodes	Returns every node_type value inside tree, keyed by its path.
paths	Returns tree's leaves keyed by their dotted path.

Transform

map	Builds a new tree with each leaf replaced by what f returns.
update	Writes path-keyed values into tree, in place.

Write your own walk

as_predicate	Resolves a Selector into a predicate.
extend_path	Extends a dotted path by one key.
flatten_one_level	Takes one interior node apart, one level deep.
is_node	Returns whether value is an interior node rather than a leaf.

Structure

TreeDef	The shape of a tree, with its leaves abstracted away.

Type aliases

Selector	a type, a tuple of types, or a predicate.