flex_model.traverse.flatten

flex_model.traverse.flatten(root_obj: Any) → Tuple[InternalNode | LeafNode | Any, List[Tensor | None]]

Flatten an arbitrary python object into a tree definition and a collection of leaves. These can then be repacked by unflatten to perfectly reconstruct the original python object. The python object is recursively unpacked using node representations, which each locally know how to unpack themselves.

Note:

The traversal is done in a depth-first way to bias us towards finding the left-most leaf node first.

Parameters:

root_obj (Any) – The python object to flatten.

Returns:

A tree definition of the python object and a list of leaf objects (typically Pytorch tensors).

Return type:

Tuple[Union[InternalNode, LeafNode, ScalarNode], List[Optional[Tensor]]]