[pytree] support `X | Y` union type in `tree_map_only` (#120389)
Follow-up PR for #119974 with some small tweaks.
1. Support `X | Y` union type for Python 3.10+
2. Enable predicate function in `tree_map_only` in CXX pytree.
3. Remove unnecessary function definition.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120389
Approved by: https://github.com/zou3519
