[pytree] add function `tree_iter` (#120155)
Fixes #119768
- #119768
This PR adds a new function `tree_iter` that lazily iterates over the tree leaves. It is different than the `tree_leaves` function while the latter traversal the whole tree first to build a list of leaves.
```python
for leaf in tree_iter(tree):
...
```
is much more efficient than:
```python
for leaf in tree_leaves(tree):
...
```
where `tree_leaves(tree)` is `list(tree_iter(tree))`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120155
Approved by: https://github.com/vmoens
