[functorch] [pytree] register return_types (pytorch/functorch#350)
* [pytree] register return_types
* add comment for binding the class in lambda
* remove tree_flatten_hack and update
* add tests
* make flake8 happy
* retrigger CI
* address review
* remove incorrectly added test
* add namedtuple test