tree_flatten_exactly_one_level

penzai.core.tree_util.tree_flatten_exactly_one_level(tree: Any) → tuple[list[tuple[Any, Any]], PyTreeDef] | None

Flattens a PyTree exactly one level, or returns None if it’s not a PyTree.

Parameters:

tree – Tree to flatten.

Returns:

If tree has any children, returns a tuple (children, treedef) where children is a list of (key, child) pairs. Otherwise, returns None.