This is an API provided by dm-tree that is not available in JAX. If we add it to jax.tree_util, it will help users who want to migrate away from dm-tree https://tree.readthedocs.io/en/latest/api.html#tree.traverse

cc/ @mtthss