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))
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.
  
trax/optimizers/base.py