code: def l2_norm(tree):
      """Returns an L2 norm computed over all elements of all tensors in `tree`.
  
      Args:
        tree: Tree-structured collection of tensors, e.g., model weights matching
            the model's layer structure.
  
      Returns:
        A scalar value computed as if all the tensors in `tree` were combined
        and flattened into a single vector, and then the L2 norm of that vector
        was calculated.
      """
      leaves = fastmath.tree_flatten(tree)
      return jnp.sqrt(sum(jnp.vdot(x, x) for x in leaves))
  
file_path: trax/trax/optimizers/base.py
imports: from trax import fastmath
  from trax.fastmath import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: l2_norm(tree) and returns: jnp