Traverse: Pytree traversal utility

Ops

flatten

Flatten an arbitrary python object into a tree definition and a collection of leaves.

unflatten

Repack a tree definition and list of leaves into the original python object.

Nodes

InternalObject

Special type indicating an unconstrained type.

LeafObject

Special type indicating an unconstrained type.

ScalarObject

Special type indicating an unconstrained type.

InternalNode

Node correponding to unpackable container.

LeafNode

Leaf node, typically corresponding to a tensor.

ScalarNode

Special type indicating an unconstrained type.