Custom PyTree Nodes not Supported #1
Description
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