code: def l2_regularization(params: PyTree):
      """Calculate the L2 loss (square L2 norm), given parameters of the model.
  
      Args:
        params: Parameters of the model.
  
      Returns:
        L2 norm.
  
      """
      weight_penalty_params = jax.tree_util.tree_leaves(params)
      return sum(
          [jnp.sum(x**2) for x in weight_penalty_params if x.ndim > 1]
      )
  
imports: import functools
  from typing import Optional, Any, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 14 for calling a global function with example usage: l2_regularization(params) and returns: sum