Add *_only and all/any pytree utilities (#83316)
With a sample usage in proxy tensor to show how they can shorten
your code dramatically.
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83316
Approved by: https://github.com/zou3519, https://github.com/albanD, https://github.com/bdhirsh