--- title: Structure inputs and outputs as pytrees description: "Flatten and rebuild nested containers of tensors, name every leaf by its key path, register your own classes and dataclasses as nodes, serialize a structure, and pass whole trees through compile, trace, eval and save." --- A pytree is a nested arrangement of Python containers whose ends are the values you care about: a dict of lists of tensors, a tuple of dicts, a dataclass holding both. `clika_runtime.pytree` takes such a tree apart into a flat list of leaves plus a `TreeSpec` that remembers the shape, and puts the two halves back together. Every boundary of the Python surface that accepts several tensors at once (`compile`, `trace`, `eval`, `save` and `load`) accepts a pytree, so a model's inputs and outputs keep the structure the code was written in. The functions carry PyTorch's spellings (`tree_flatten`, `tree_map`, `register_pytree_node`), so code written against `torch.utils._pytree` reads the same here. Two rules differ from a plain container walk and are worth reading first. ## Leaves, and why None is one of them Lists, tuples, dicts, `OrderedDict`, `defaultdict`, `deque`, namedtuples, and every class registered with the node registry are nodes; everything else (numbers, strings, bytes, tensors, objects the registry does not know) is a leaf. Leaves come out depth first, left to right, and a flatten followed by an unflatten reproduces the same container types. ```python title="flatten.py" from clika_runtime import pytree tree = {"a": [1, (2, None)], "b": 3, "c": {"d": (4,), "e": []}} leaves, spec = pytree.tree_flatten(tree) print(leaves) # [1, 2, None, 3, 4] print(spec.num_leaves) # 5 print(spec) # TreeSpec({'a': [*, (*, *)], 'b': *, 'c': {'d': (*,), 'e': []}}) rebuilt = pytree.tree_unflatten(spec, leaves) assert rebuilt == tree and type(rebuilt) is dict ``` `None` is a leaf. A model that takes an optional input (a KV cache that is empty on the first step) keeps the same flat index for every other input whether or not the optional one is present, so a compiled graph's argument slots stay stable. `none_is_leaf=False` turns `None` into a childless node instead, and `is_leaf=lambda x: isinstance(x, list)` stops the descent at every list, which then counts as one leaf. ```python title="leaf_rules.py" from clika_runtime import pytree with_cache = {"input_ids": 1, "cache": (2, 3)} without_cache = {"input_ids": 1, "cache": None} leaves_a, _ = pytree.tree_flatten(with_cache) leaves_b, spec_b = pytree.tree_flatten(without_cache) print(leaves_a[0], leaves_b[0]) # 1 1 print(leaves_b) # [1, None] print(pytree.tree_unflatten(spec_b, [10, None])) # {'input_ids': 10, 'cache': None} leaves, spec = pytree.tree_flatten({"x": 1, "cache": None}, none_is_leaf=False) print(leaves) # [1] print(spec) # TreeSpec({'x': *, 'cache': None}, none_is_leaf=False) ``` ## Map and reduce `tree_map` applies a function to every leaf and rebuilds the same structure; extra trees pair their leaves by position and must share the first tree's structure (a mismatch raises `ValueError`). `tree_map_` runs the function for its side effect and returns the tree it was given; `tree_map_only` filters by type or by predicate. The reductions (`tree_all`, `tree_any`, `tree_reduce`, `tree_sum`, `tree_max`, `tree_min`) read the leaves without rebuilding anything, `tree_leaves` returns the flat list and `tree_iter` yields it lazily. ```python title="map.py" from clika_runtime import pytree print(pytree.tree_map(lambda x: x is None, {"x": 1, "y": None})) # {'x': False, 'y': True} print(pytree.tree_map(lambda x, y: x + y, {"a": 1, "b": 2}, {"a": 10, "b": 20})) # {'a': 11, 'b': 22} mixed = {"a": 1, "b": "text", "c": [2, None, 2.5]} print(pytree.tree_map_only(int, lambda x: x + 1, mixed)) # {'a': 2, 'b': 'text', 'c': [3, None, 2.5]} ``` ## Name every leaf by its key path `tree_flatten_with_path` returns `(path, leaf)` pairs; a path is a tuple of key entries (`MappingKey`, `SequenceKey`, `GetAttrKey`, `DataclassKey`), `keystr` renders it as the indexing expression that reaches the leaf, and `key_get` follows it. Key paths are how the file format names the leaves of a saved tree (below) and how a diagnostic points at one tensor inside a large state. ```python title="paths.py" from typing import NamedTuple from clika_runtime import pytree class Point(NamedTuple): x: object y: object tree = {"a": [1, 2], "p": Point(3, 4)} pairs = pytree.tree_flatten_with_path(tree)[0] print([leaf for _, leaf in pairs]) # [1, 2, 3, 4] print([pytree.keystr(path) for path, _ in pairs]) # ["['a'][0]", "['a'][1]", "['p'].x", "['p'].y"] path = (pytree.MappingKey("a"), pytree.SequenceKey(1)) print(pytree.key_get(tree, path)) # 2 ``` ## Register your own containers A class the registry does not know is a leaf. `register_pytree_node(cls, flatten_fn, unflatten_fn)` opens it: `flatten_fn` returns `(children, context)` or `(children, context, entries)` (the entries name the children for key paths), and `unflatten_fn(children, context)` rebuilds an instance. The argument order is PyTorch's. After registration, every tree function descends into the class and every rebuild produces an instance of it. ```python title="register.py" from clika_runtime import pytree class Point: def __init__(self, x, y): self.x, self.y = x, y pytree.register_pytree_node( Point, lambda p: ((p.x, p.y), None, ("x", "y")), lambda children, _ctx: Point(*children), ) print(pytree.tree_leaves({"p": Point(1, 2), "n": 3})) # [1, 2, 3] mapped = pytree.tree_map(lambda v: v * 10, Point(1, 2)) print(type(mapped).__name__, mapped.x, mapped.y) # Point 10 20 spec = pytree.tree_structure(Point(1, 2)) print(spec.entries(), spec.paths()) # ['x', 'y'] [('x',), ('y',)] ``` A class can carry its own flattening as a method pair and register with the decorator. Note the method order: `__tree_unflatten__(cls, context, children)` takes the context first, while the function form above takes `(children, context)`. ```python title="register_class.py" from clika_runtime import pytree @pytree.register_pytree_node_class class Pair: def __init__(self, a, b): self.a, self.b = a, b def __tree_flatten__(self): return (self.a, self.b), "pair", ("a", "b") @classmethod def __tree_unflatten__(cls, context, children): return cls(*children) leaves, spec = pytree.tree_flatten(Pair(1, [2, 3])) print(leaves, spec.context) # [1, 2, 3] pair rebuilt = pytree.tree_unflatten(spec, [10, 20, 30]) print(rebuilt.a, rebuilt.b) # 10 [20, 30] ``` Dataclasses register by field: the data fields become children, the meta fields ride the context and come back unchanged. ```python title="register_dataclass.py" import dataclasses from clika_runtime import pytree @dataclasses.dataclass class Sample: weight: object bias: object name: str pytree.register_dataclass(Sample, data_fields=["weight", "bias"], meta_fields=["name"]) leaves, spec = pytree.tree_flatten(Sample(1, 2, "s")) print(leaves) # [1, 2] print(pytree.tree_unflatten(spec, [10, 20])) # Sample(weight=10, bias=20, name='s') print([pytree.keystr(p) for p, _ in pytree.tree_flatten_with_path(Sample(1, 2, "s"))[0]]) # ['.weight', '.bias'] ``` Registering a built-in container, an instance instead of a class, or the same class twice raises; `unregister_pytree_node(cls)` makes instances leaves again. A registration meant for one library goes into a namespace so it never changes what other code sees: `register_pytree_node(..., namespace="mylib")` opens the class only for calls that pass `namespace="mylib"` (`tree_leaves(tree, namespace="mylib")`), `is_registered(cls, namespace="mylib")` reports it, and a namespace with no registration of its own falls back to the global one. ## Dictionary order Dicts flatten in insertion order, and the order is part of the structure: `{"b": 1, "a": 2}` and `{"a": 2, "b": 1}` have different specs. Code that wants two dicts with the same keys to share one structure regardless of insertion order wraps the calls in `dict_insertion_ordered(False)`, which flattens by sorted key inside the block; the rebuilt dict still comes back in the original order. ```python title="dict_order.py" from clika_runtime import pytree tree = {"b": 1, "a": 2, "c": 3} leaves, spec = pytree.tree_flatten(tree) print(leaves) # [1, 2, 3] print(list(pytree.tree_unflatten(spec, leaves))) # ['b', 'a', 'c'] with pytree.dict_insertion_ordered(False): print(pytree.tree_leaves({"b": 1, "a": 2})) # [2, 1] print(pytree.tree_structure({"b": 1, "a": 2}) == pytree.tree_structure({"a": 0, "b": 0})) # True ``` ## Treespecs A `TreeSpec` is a value: it compares and hashes by structure, prints as the container shape with a `*` per leaf, and rebuilds a tree from any sequence of the right length (`spec.unflatten(leaves)`; a wrong count raises `ValueError` naming the expected number). `is_prefix` asks whether one structure is the top of another, `compose` grafts an inner structure onto every leaf of an outer one, and `flatten_up_to` flattens a tree only as deep as the spec goes. ```python title="treespec.py" from clika_runtime import pytree spec = pytree.tree_structure({"a": [1, (2, None)], "b": 3}) print(spec) # TreeSpec({'a': [*, (*, *)], 'b': *}) print(spec.num_leaves, spec.num_nodes, spec.num_children) # 4 7 2 print(spec.paths()) # [('a', 0), ('a', 1, 0), ('a', 1, 1), ('b',)] short = pytree.tree_structure([1, 2]) deep = pytree.tree_structure([1, (2, 3)]) print(short.is_prefix(deep), deep.is_prefix(short)) # True False outer = pytree.tree_structure([1, 2]) inner = pytree.tree_structure((1, 2)) print(outer.compose(inner)) # TreeSpec([(*, *), (*, *)]) ``` `treespec_dumps` writes a spec as JSON and `treespec_loads` reads it back, so a structure can travel beside a file or a request. Built-in containers and namedtuples serialize as they are; a registered class needs a `serialized_type_name`, and a context that is not JSON needs the `to_dumpable_context` / `from_dumpable_context` pair at registration. An unnamed custom node or an unknown type name in the document raises `ValueError`. ```python title="serialize.py" import json from clika_runtime import pytree spec = pytree.tree_structure({"a": [1, (2, None)], "b": 3}) text = pytree.treespec_dumps(spec) print(json.loads(text)["version"] == pytree.SERIALIZATION_PROTOCOL) # True loaded = pytree.treespec_loads(text) print(loaded == spec) # True print(loaded.unflatten([1, 2, None, 3])) # {'a': [1, (2, None)], 'b': 3} ``` ## Pytrees at the boundaries `crt.compile` takes a callable whose arguments and return value are pytrees of tensors: it flattens the call's inputs, treats the tensor leaves as graph inputs and every other leaf as a static value that is part of the guard, and rebuilds the output structure on the way out; `dynamic=True` keeps one graph across batch sizes. ```python title="compile_tree.py" import numpy as np import clika_runtime as crt import clika_runtime.nn as nn BATCH, D_IN, D_OUT = 5, 131, 64 rng = np.random.default_rng(0) w = rng.standard_normal((D_IN, D_OUT)).astype(np.float32) * 0.1 b = rng.standard_normal(D_OUT).astype(np.float32) x = rng.standard_normal((BATCH, D_IN)).astype(np.float32) class Head(nn.Module): def __init__(self, w: np.ndarray, b: np.ndarray) -> None: super().__init__() self.weight = nn.Parameter(crt.tensor(w)) self.bias = nn.Parameter(crt.tensor(b)) def forward(self, x: crt.Tensor) -> crt.Tensor: return crt.softmax(crt.relu(crt.add(crt.matmul(x, self.weight), self.bias)), dim=-1) class Wrapped(nn.Module): def __init__(self) -> None: super().__init__() self.head = Head(w, b) def forward(self, batch: dict[str, crt.Tensor]) -> dict[str, crt.Tensor]: p = self.head(batch["x"]) return {"probs": p, "argmax": crt.argmax(p, dims=[1])} compiled = crt.compile(Wrapped(), dynamic=True) out = compiled({"x": crt.tensor(x)}) print(sorted(out)) # ['argmax', 'probs'] again = compiled({"x": crt.tensor(np.tile(x, (2, 1)))}) print(tuple(again["probs"].shape)) # (10, 64) ``` `crt.trace` records a function once over stand-ins and returns a graph that runs by position or by name; the traced function takes its tensor inputs as one list and returns a list, and [Trace eager code to graphs](trace-eager-code-to-graphs.mdx) walks through it. Operators return before their work runs. `crt.eval(*trees)` flattens whatever trees it is given, skips the leaves that are not tensors, and returns once every tensor leaf holds its value; a later read then costs no wait. ```python title="eval_tree.py" import numpy as np import clika_runtime as crt x = np.random.default_rng(0).standard_normal((5, 131)).astype(np.float32) t = crt.tensor(x) tree = {"a": crt.exp(t), "b": [crt.sum(t), None, "text"], "c": (crt.relu(t),)} crt.eval(tree, crt.abs(t)) print(np.allclose(tree["a"].numpy(), np.exp(x.astype(np.float64)), rtol=1e-5, atol=1e-6)) # True ``` `crt.save` writes a tensor, a state dict, or any pytree of tensors as a safetensors file, and `crt.load` reads it back. A state dict (`crt.save(state, path)`, `crt.load(path, device="cpu")`) saves under its own keys in insertion order, so another safetensors reader sees it as it is; a deeper tree saves one entry per leaf named by `keystr` of its key path, with the structure carried in the file's metadata, and `crt.load` rebuilds the same tree. ```python title="save_tree.py" from pathlib import Path import numpy as np import clika_runtime as crt from clika_runtime import pytree rng = np.random.default_rng(0) tree = { "layers": [ {"w": crt.tensor(rng.standard_normal((32, 131)).astype(np.float32)), "b": crt.tensor(np.zeros(32, np.float32))} for _ in range(2) ], "step": crt.tensor(np.array([7], dtype=np.int64)), } crt.save(tree, Path("tree.safetensors")) back = crt.load(Path("tree.safetensors")) print(pytree.tree_structure(back) == pytree.tree_structure(tree)) # True print(sorted(pytree.keystr(kp) for kp, _ in pytree.tree_flatten_with_path(tree)[0])) # ["['layers'][0]['b']", "['layers'][0]['w']", "['layers'][1]['b']", "['layers'][1]['w']", "['step']"] ``` The same structure rule holds for a module's state: `crt.save(module.state_dict(), path)` followed by `other.load_state_dict(crt.load(path))` reproduces the forward, and the file is plain safetensors any other tool reads. [Use ClikaRT from Python](use-clikart-from-python.mdx) covers the tensor and module surface the leaves belong to.