Tree API

Nested-structure helpers with the public API of jax.tree. Pure Python; the only binary dependency is NumPy (no JAX, no C++ pytree). None is an empty pytree.

Public functions and built-in walking rules are stable. Custom-type registration (__tree_flatten__ / __tree_unflatten__ and JAX / Torch / optree registries) is beta — see Overview. Why this exists (and how it relates to jax.tree / dm-tree / optree): Why tree.

anytensor.tree

Nested-structure utilities with the public API of jax.tree.

Pure Python; the only binary dependency is NumPy (no JAX, no C++ pytree extension). Walking rules follow jax.tree / jax.tree_util:

• None is an empty pytree (zero leaves), not a leaf.
• flatten(tree) returns (leaves, treedef).
• Dicts flatten by sorted keys; OrderedDict keeps insertion order.
• Arrays / tensors are leaves. str / bytes / sets / mapping views are leaves.

Custom nodes

1. Magic flatten (beta)::

   def tree_flatten(self): return children, aux

   @classmethod def tree_unflatten(cls, aux, children): return cls(...)

2. Batch / unbatch (stable; same functions as :mod:anytensor.jraph batch / unbatch). Checked before walking children. Use AnyTensor ops so the type stays portable::

   @classmethod def tree_batch(cls, xs, axis=0): return cls(at.concatenate([x.values for x in xs], axis=axis))

   def tree_unbatch(self, axis=0): n = int(at.shape(self.values)[axis]) ids = at.arange(n, like=self.values) return [ cls(at.take(self.values, ids[i : i + 1], axis=axis)) for i in range(n) ]

GraphsTuple implements these (offset senders/receivers). Objects without magic unbatch along the leading axis into unit slices.

1. Already-imported pytree registries (beta), looked up by type (never imported as a side effect): jax.tree_util, torch.utils._pytree, and optree. Built-in containers stay on this module's path.

Public map / flatten / batch / unbatch and built-in walking rules are stable. Flatten-style registration (item 1 and item 3) may change.

SequenceKey

Bases: NamedTuple

Path entry for a list/tuple child (jax.tree name).

Source code in anytensor/tree.py

class SequenceKey(NamedTuple):
    """Path entry for a list/tuple child (``jax.tree`` name)."""

    idx: int

DictKey

Bases: NamedTuple

Path entry for a mapping child.

Source code in anytensor/tree.py

class DictKey(NamedTuple):
    """Path entry for a mapping child."""

    key: Any

GetAttrKey

Bases: NamedTuple

Path entry for a namedtuple / attrs field.

Source code in anytensor/tree.py

class GetAttrKey(NamedTuple):
    """Path entry for a namedtuple / attrs field."""

    name: str

PyTreeDef

Tree structure leftover after flattening (jax.tree_util.PyTreeDef-like).

Source code in anytensor/tree.py

class PyTreeDef:
    """Tree structure leftover after flattening (``jax.tree_util.PyTreeDef``-like)."""

    __slots__ = ("kind", "metadata", "children", "_restore")

    def __init__(self, kind: str, metadata, children, restore=None):
        self.kind = kind
        self.metadata = metadata
        self.children = tuple(children)
        self._restore = restore

    @property
    def num_leaves(self) -> int:
        if self.kind == "leaf":
            return 1
        return sum(c.num_leaves for c in self.children)

    def __eq__(self, other):
        if type(other) is not PyTreeDef:
            return NotImplemented
        return (
            self.kind == other.kind
            and self.metadata == other.metadata
            and self.children == other.children
        )

    def __repr__(self) -> str:
        if self.kind == "leaf":
            return "PyTreeDef(*)"
        if self.kind == "none":
            return "PyTreeDef(None)"
        kids = ", ".join(repr(c)[len("PyTreeDef") :] if repr(c).startswith("PyTreeDef") else repr(c) for c in self.children)
        # Keep repr short and stable; exact JAX text is not required.
        return f"PyTreeDef({self.kind}[{kids}])"

    def unflatten(self, leaves):
        leaves = list(leaves)
        out, i = self._unflatten(leaves, 0)
        if i != len(leaves):
            raise ValueError(
                f"Too many leaves for PyTreeDef; expected {self.num_leaves}, got {len(leaves)}"
            )
        return out

    def _unflatten(self, leaves, i):
        if self.kind == "leaf":
            if i >= len(leaves):
                raise ValueError(
                    f"Too few leaves for PyTreeDef; expected {self.num_leaves}, got {len(leaves)}"
                )
            return leaves[i], i + 1
        if self.kind == "none":
            return None, i
        vals = []
        for child in self.children:
            v, i = child._unflatten(leaves, i)
            vals.append(v)
        return _rebuild(self, vals), i

    def flatten_up_to(self, tree) -> list:
        """Flatten ``tree`` using this schema; stop at each leaf of ``self``.

        ``self`` must be a prefix of ``tree``'s structure (same as JAX). Values
        at leaf positions are returned as-is — so a post-``split`` tree of
        chunk lists can be flattened with the unsplit batch's treedef.
        """
        acc: list = []
        self._flatten_up_to(tree, acc)
        return acc

    def _flatten_up_to(self, node, acc) -> None:
        if self.kind == "leaf":
            acc.append(node)
            return
        if self.kind == "none":
            if node is not None:
                raise ValueError(
                    "pytree structure error: trees must have the same structure."
                )
            return
        entry = _one_level(node)
        if entry is None or entry[0] != self.kind or entry[1] != self.metadata:
            raise ValueError(
                "pytree structure error: trees must have the same structure."
            )
        _kind, _metadata, children, _restore = entry
        if len(children) != len(self.children):
            raise ValueError(
                "pytree structure error: trees must have the same structure."
            )
        for child_def, child in zip(self.children, children):
            child_def._flatten_up_to(child, acc)

flatten_up_to

flatten_up_to(tree)

Flatten tree using this schema; stop at each leaf of self.

self must be a prefix of tree's structure (same as JAX). Values at leaf positions are returned as-is — so a post-split tree of chunk lists can be flattened with the unsplit batch's treedef.

Source code in anytensor/tree.py

def flatten_up_to(self, tree) -> list:
    """Flatten ``tree`` using this schema; stop at each leaf of ``self``.

    ``self`` must be a prefix of ``tree``'s structure (same as JAX). Values
    at leaf positions are returned as-is — so a post-``split`` tree of
    chunk lists can be flattened with the unsplit batch's treedef.
    """
    acc: list = []
    self._flatten_up_to(tree, acc)
    return acc

flatten

flatten(tree, is_leaf=None)

Flatten tree into (leaves, treedef).

import anytensor.tree as tree leaves, _ = tree.flatten((1, (2, 3))) tuple(leaves) (1, 2, 3) empty, _ = tree.flatten(None) len(empty) 0

Source code in anytensor/tree.py

def flatten(tree, is_leaf=None):
    """Flatten ``tree`` into ``(leaves, treedef)``.

    >>> import anytensor.tree as tree
    >>> leaves, _ = tree.flatten((1, (2, 3)))
    >>> tuple(leaves)
    (1, 2, 3)
    >>> empty, _ = tree.flatten(None)
    >>> len(empty)
    0
    """
    acc: list = []
    treedef = _flatten_into(tree, acc, is_leaf)
    return acc, treedef

unflatten

unflatten(treedef, leaves)

Rebuild a tree from treedef and leaves.

Source code in anytensor/tree.py

def unflatten(treedef: PyTreeDef, leaves: Iterable):
    """Rebuild a tree from ``treedef`` and ``leaves``."""
    return treedef.unflatten(leaves)

leaves

leaves(tree, is_leaf=None)

Return the leaves of tree (None yields []).

Source code in anytensor/tree.py

def leaves(tree, is_leaf=None) -> list:
    """Return the leaves of ``tree`` (``None`` yields ``[]``)."""
    return flatten(tree, is_leaf=is_leaf)[0]

structure

structure(tree, is_leaf=None)

Return the PyTreeDef of tree.

Source code in anytensor/tree.py

def structure(tree, is_leaf=None) -> PyTreeDef:
    """Return the ``PyTreeDef`` of ``tree``."""
    return flatten(tree, is_leaf=is_leaf)[1]

map

map(f, tree, *rest, is_leaf=None)

Map f over the leaves of tree (and rest).

Subsequent trees are flattened with tree's schema via :meth:PyTreeDef.flatten_up_to (JAX-style prefix).

import anytensor.tree as tree tree.map(lambda v: v * 2, {"b": 1, "a": [2, 3]}) {'a': [4, 6], 'b': 2} tree.map(lambda x: x + 1, None) is None True

Source code in anytensor/tree.py

def map(f, tree, *rest, is_leaf=None):  # noqa: A001
    """Map ``f`` over the leaves of ``tree`` (and ``rest``).

    Subsequent trees are flattened with ``tree``'s schema via
    :meth:`PyTreeDef.flatten_up_to` (JAX-style prefix).

    >>> import anytensor.tree as tree
    >>> tree.map(lambda v: v * 2, {"b": 1, "a": [2, 3]})
    {'a': [4, 6], 'b': 2}
    >>> tree.map(lambda x: x + 1, None) is None
    True
    """
    leaves0, treedef = flatten(tree, is_leaf=is_leaf)
    rest_leaves = [treedef.flatten_up_to(other) for other in rest]
    if rest_leaves:
        out = [f(*xs) for xs in zip(leaves0, *rest_leaves)]
    else:
        out = [f(x) for x in leaves0]
    return unflatten(treedef, out)

flatten_with_path

flatten_with_path(tree, is_leaf=None)

Flatten into ([(path, leaf), ...], treedef).

Source code in anytensor/tree.py

def flatten_with_path(tree, is_leaf=None):
    """Flatten into ``([(path, leaf), ...], treedef)``."""
    acc: list = []
    treedef = _flatten_with_path_into(tree, (), acc, is_leaf)
    return acc, treedef

leaves_with_path

leaves_with_path(tree, is_leaf=None)

Return [(path, leaf), ...].

Source code in anytensor/tree.py

def leaves_with_path(tree, is_leaf=None) -> list:
    """Return ``[(path, leaf), ...]``."""
    return flatten_with_path(tree, is_leaf=is_leaf)[0]

map_with_path

map_with_path(f, tree, *rest, is_leaf=None)

Like :func:map but f receives (path, *leaves).

Source code in anytensor/tree.py

def map_with_path(f, tree, *rest, is_leaf=None):
    """Like :func:`map` but ``f`` receives ``(path, *leaves)``."""
    pairs, treedef = flatten_with_path(tree, is_leaf=is_leaf)
    paths = [p for p, _ in pairs]
    leaves0 = [v for _, v in pairs]
    rest_leaves = [treedef.flatten_up_to(other) for other in rest]
    if rest_leaves:
        out = [f(p, *xs) for p, xs in zip(paths, zip(leaves0, *rest_leaves))]
    else:
        out = [f(p, x) for p, x in zip(paths, leaves0)]
    return unflatten(treedef, out)

all

all(tree, *, is_leaf=None)

True if every leaf is truthy (empty trees are True).

Source code in anytensor/tree.py

def all(tree, *, is_leaf=None):  # noqa: A001
    """``True`` if every leaf is truthy (empty trees are ``True``)."""
    return builtins.all(leaves(tree, is_leaf=is_leaf))

reduce

reduce(function, tree, initializer=_UNSET, is_leaf=None)

Reduce leaves with function (same empty-tree error as functools.reduce).

Source code in anytensor/tree.py

def reduce(function, tree, initializer=_UNSET, is_leaf=None):  # noqa: A002
    """Reduce leaves with ``function`` (same empty-tree error as ``functools.reduce``)."""
    xs = leaves(tree, is_leaf=is_leaf)
    if initializer is _UNSET:
        return _f_reduce(function, xs)
    return _f_reduce(function, xs, initializer)

batch

batch(trees, axis=0)

Batch a sequence of pytrees along axis (leading axis by default).

Same function as :func:anytensor.jraph.batch. If the type defines __tree_batch__(xs, axis=0), that method is used and children are not walked. Otherwise leaves are concatenated with :func:anytensor.concatenate. All-None stays None.

import numpy as np import anytensor.tree as tree tree.batch([np.array([1, 2]), np.array([3])]) array([1, 2, 3])

Source code in anytensor/tree.py

def batch(trees, axis: int = 0):
    """Batch a sequence of pytrees along ``axis`` (leading axis by default).

    Same function as :func:`anytensor.jraph.batch`. If the type defines
    ``__tree_batch__(xs, axis=0)``, that method is used and children are
    **not** walked. Otherwise leaves are concatenated with
    :func:`anytensor.concatenate`. All-``None`` stays ``None``.

    >>> import numpy as np
    >>> import anytensor.tree as tree
    >>> tree.batch([np.array([1, 2]), np.array([3])])
    array([1, 2, 3])
    """
    xs = list(trees)
    if not xs:
        raise ValueError("batch() requires at least one structure")
    return _batch_impl(xs, axis=axis)

unbatch

unbatch(structure, axis=0)

Unbatch structure along axis into unit slices.

Same function as :func:anytensor.jraph.unbatch. __tree_unbatch__(axis=0) on the object wins (GraphsTuple yields one graph per n_node entry). Nested containers recurse. Top-level None cannot infer a batch size.

import numpy as np import anytensor.tree as tree head, tail = tree.unbatch(np.arange(1, 3)) tuple(int(x) for x in head), tuple(int(x) for x in tail) ((1,), (2,))

Source code in anytensor/tree.py

def unbatch(structure, axis: int = 0):
    """Unbatch ``structure`` along ``axis`` into unit slices.

    Same function as :func:`anytensor.jraph.unbatch`. ``__tree_unbatch__(axis=0)``
    on the object wins (``GraphsTuple`` yields one graph per ``n_node`` entry).
    Nested containers recurse. Top-level ``None`` cannot infer a batch size.

    >>> import numpy as np
    >>> import anytensor.tree as tree
    >>> head, tail = tree.unbatch(np.arange(1, 3))
    >>> tuple(int(x) for x in head), tuple(int(x) for x in tail)
    ((1,), (2,))
    """
    if structure is not None and hasattr(type(structure), "__tree_unbatch__"):
        return list(_call_unbatch(structure, axis))
    if structure is None:
        raise ValueError("Cannot unbatch None")
    n = _leading_length(structure, axis)
    return _unbatch_units(structure, n, axis)