google
Deprecate/Remove traverse_util module
As per @marcvanzee's comment on #1022, most of the traverse_util logic can be done with tree_maps without losing any run time optimizations. The trade-off is the compile times will be longer, but tree_map code is easier to understand.
cc: @avital
I think flatten_dict and unflatten_dict are helpful utilities, perhaps they can be moved to jax_utils?
I agree, those are useful. I don't know what the long term best place for them is, though. Maybe even outside of Flax? So I guess maybe it's fine to just keep them where they are for now (jax_utils is a bad package -- it doesn't describe what it's scope is... traverse_utils is almost descriptive in this case.) But I don't have strong feelings here.
Just curious @n2cholas -- are you planning to work on this or did you file it for us to keep track of it (both answers are fine, just checking). If you don't want to work on this I'll add "pull requests encouraged" to it and maybe someone else wants to give it a shot.
I'd be happy to work on this, but what's the process to actually deprecate then remove this module? Do I just add deprecation warnings to all the functions?
As a first step, yes. Add a deprecation warning and make sure it is clearly deprecated in the docs. Make sure we don't use them anywhere in our examples/HOWTOs/etc. Then we can cut a new release, and in a few months remove it entirely from our code.