Minor improvements and cleanup using tree utility functions.

[carlosgmartin]

Edits code to use already-existent tree utility functions, where possible.

[vroulet]

Make it three PRs :)
For optax.global_norm vs optax.tree.norm, just to be extra-sure (since that optax.global_norm is used a lot), can you make sure that their hlos match?
For tree.norm, keep the same signature as jnp.linalg.norm (that's what users will probably expect.)

[carlosgmartin]

@vroulet Done. Here are the other 2 PRs:

• Deprecate optax.global_norm in favor of optax.tree.norm. #1368
• Streamline interface of tree_norm function. #1369

[carlosgmartin]

@vroulet Are there any changes you'd like me to make to this PR?

[rdyro]

Thanks @carlosgmartin ! Let's merge this!