Minor improvements and cleanup using tree utility functions.#1365

copybara-service[bot] merged 1 commit intogoogle-deepmind:maingoogle-deepmind/optax:mainfrom

carlosgmartin:minor_improvements_using_tree_utilscarlosgmartin/optax:minor_improvements_using_tree_utilsCopy head branch name to clipboard