Add jax.tree_util.tree_traverse

[jakevdp]

This provides the same API as tree.traverse from https://github.com/google-deepmind/tree

cc/ @mtthss

Fixes #19631

[XuehaiPan]

It would be better to use itertools.count to handle self-referential cases.

import itertools

counter = itertools.count()
one_level_children, one_level_treedef = tree_flatten(structure, is_leaf=lambda x: next(counter) > 0)

In [1]: from jax.tree_util import tree_flatten

In [2]: d = {'a': 1, 'b': 2, 'c': 3}; d['d'] = d

In [3]: tree_flatten(d, is_leaf=lambda x: x is not d)
RecursionError: maximum recursion depth exceeded while calling a Python object

In [4]: import itertools

In [5]: counter = itertools.count()

In [6]: tree_flatten(d, is_leaf=lambda x: next(counter) > 0)
Out[6]: ([1, 2, 3, {'a': 1, 'b': 2, 'c': 3, 'd': ...}], PyTreeDef({'a': *, 'b': *, 'c': *, 'd': *}))

[superbobry]

I don't think we should handle self-referential cases, since the underlying structure is not a tree.

[ayaka14732]

@superbobry I think there is no need to add __str__ = __repr__ here, because str() can use __repr__ when __str__ is not available.

From Python documentation:

If object does not have a str() method, then str() falls back to returning repr(object).

[superbobry]

Good call, thanks for pointing this out, Ayaka!

[ayaka14732]

Maybe need to avoid double instantiation and subclassing, so that there is only one instance of class _MapToNone.

I've implemented similar things before: https://github.com/ayaka14732/einshard/blob/00277824cbf75eb2b01d5cffb1935ab7c6fdc1c6/src/einshard/parsec/Void.py

[ayaka14732]

Maybe export class _MapToNone as well since it may be used in writting type signatures.