def tree_flatten(tree, is_leaf: Optional[Callable[[Any], bool]] = None):
"""Flattens a pytree.
Args:
tree: a pytree to flatten.
is_leaf: an optionally specified function that will be called at each
flattening step. It should return a boolean, with true stopping the
traversal and the whole subtree being treated as a leaf, and false
indicating the flattening should traverse the current object.
Returns:
A pair where the first element is a list of leaf values and the second
element is a treedef representing the structure of the flattened tree.
"""
return pytree.flatten(tree, is_leaf)
def tree_structure(tree, is_leaf: Optional[Callable[[Any], bool]] = None): """Gets the treedef for a pytree.""" return pytree.flatten(tree, is_leaf)[1]
def tree_leaves(tree, is_leaf: Optional[Callable[[Any], bool]] = None): """Gets the leaves of a pytree.""" return pytree.flatten(tree, is_leaf)[0]
def _process_pytree(process_node, tree): leaves, treedef = pytree.flatten(tree) return treedef.walk(process_node, None, leaves), treedef
def tree_structure(tree): """Gets the treedef for a pytree.""" return pytree.flatten(tree)[1]
def tree_leaves(tree): """Gets the leaves of a pytree.""" return pytree.flatten(tree)[0]
