tree_util

Additional tree functionality related to jax.tree_util.

Classes

CustomGetAttrKey

Subclass-friendly variant of jax.tree_util.GetAttrKey.

Functions

pretty_keystr(keypath, tree)	Constructs a pretty name from a keypath and an object.
tree_flatten_exactly_one_level(tree)	Flattens a PyTree exactly one level, or returns None if it's not a PyTree.