code: def l2_norm(tree):
      """Compute the l2 norm of a pytree of arrays. Useful for weight decay."""
      leaves, _ = tree_flatten(tree)
      return jnp.sqrt(sum(jnp.vdot(x, x) for x in leaves))
  
imports: from __future__ import annotations
  from typing import Any, Callable, NamedTuple
  from collections import namedtuple
  import functools
  from functools import partial
  import jax.numpy
  from jax._src.util import safe_zip, safe_map, unzip2
  from jax import tree_util
  from jax.tree_util import tree_map, tree_flatten, tree_unflatten, register_pytree_node
label_desc: Use image node_id 17 for calling a global function with example usage: l2_norm(tree) and returns: jnp