Streamline interface of tree_norm function.#1369

carlosgmartin wants to merge 1 commit intogoogle-deepmind:maingoogle-deepmind/optax:mainfrom

carlosgmartin:streamline_tree_normcarlosgmartin/optax:streamline_tree_normCopy head branch name to clipboard