all(tree, *[, is_leaf])

Call all() over the leaves of a tree.

broadcast(prefix_tree, full_tree[, is_leaf])

Broadcasts a tree prefix into the full structure of a given tree.

flatten(tree[, is_leaf])

Flattens a pytree.

flatten_with_path(tree[, is_leaf, ...])

Flattens a pytree liketree_flatten, but also returns each leaf's key path.

leaves(tree[, is_leaf])

Gets the leaves of a pytree.

leaves_with_path(tree[, is_leaf, ...])

Gets the leaves of a pytree liketree_leaves and returns each leaf's key path.

map(f, tree, *rest[, is_leaf])

Maps a multi-input function over pytree args to produce a new pytree.

map_with_path(f, tree, *rest[, is_leaf, ...])

Maps a multi-input function over pytree key path and args to produce a new pytree.

reduce(function, tree[, initializer, is_leaf])

Call reduce() over the leaves of a tree.

reduce_associative(operation, tree, *[, ...])

Perform a reduction over a pytree with an associative binary operation.

static(**kwargs)

Convenience wrapper to declare a static pytree attribute.

structure(tree[, is_leaf])

Gets the treedef for a pytree.

transpose(outer_treedef, inner_treedef, ...)

Transform a tree having tree structure (outer, inner) into one having structure (inner, outer).

unflatten(treedef, leaves)

Reconstructs a pytree from the treedef and the leaves.