etree

[[Source]]

from etils import etree

Tree utils.

All symbols

Module

etree	Tree utils.
etree.backend	Tree API backends.
etree.tree_utils	Tree API.

Class

etree.backend.Backend	Tree API backend.
etree.backend.DmTree	tree backend.
etree.backend.Jax	jax.tree_util backend.
etree.backend.Nest	tf.nest backend.
etree.backend.Optree	optree backend.
etree.backend.Python	Pure Python backend.
etree.tree_utils.TreeAPI	Tree API, using either jax.tree_utils, tf.nest or tree backend.

Function

etree.map	Same as tree.map_structure.
etree.parallel_map	Same as tree.map_structure but apply map_fn in parallel.
etree.spec_like	Inspect a tree of array, works with any array type.
etree.stack	Stack a tree of Iterable[Array].
etree.unzip	Unpack a tree of iterable.

Attribute

etree.jax	Expose 4 variants of the API depending on which backend is used.
etree.nest
etree.optree
etree.py
etree.tree

Typing

etree.Tree	TODO(b/202712189): Add Any to the Union Tree can be arbitrary `chex.dataclas…