tree_util

Additional tree functionality related to jax.tree_util.

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.