PyTree path aware traverse in jax

[port2077]

Hi,

I am trying to label my params leaves based on their path in the Pytree object. Is there a way to traverse this in jax in a way that is probably similar to the api traverse_util.path_aware_map from the Flax library ?

[jakevdp]

Since you closed the issue I suspect you already found this, but in case people come across this question by search engine, the relevant API is jax.tree.map_with_path: https://docs.jax.dev/en/latest/_autosummary/jax.tree.map_with_path.html