Currently, only PyTrees consisting of arbitrary nestings of Dict, Tuple, and List are supported. Strictly speaking, PyTrees can consist of any "container" type, and custom containers, like NamedTuple, can be registered.
Support for these node types should be added to the various methods in tree_util
Currently, only PyTrees consisting of arbitrary nestings of
Dict,Tuple, andListare supported. Strictly speaking, PyTrees can consist of any "container" type, and custom containers, like NamedTuple, can be registered.Support for these node types should be added to the various methods in
tree_util