code: def clip_grads(grad_tree, max_norm):
      """Proportionally reduces each gradient value to respect an aggregate limit.
  
      Args:
        grad_tree: Gradient values structured as a tree of tensors matching the
            model's layer structure.
        max_norm: The aggregate limit on gradient values. All gradient elements in
            `grad_tree` are treated as if they belonged to a single vector and
            that vector is shortened if needed so that its L2 norm does not exceed
            `clip_grad_norm`.
  
      Returns:
        A new tree of tensors matching the structure of `grad_tree`, but with
        element values proportionally rescaled as needed to respect the `max_norm`
        limit.
      """
      norm = l2_norm(grad_tree)
      normalize = lambda g: jnp.where(
          norm < max_norm, g, g * (max_norm / norm)
      )
      return fastmath.nested_map(grad_tree, normalize)
  
file_path: trax/trax/optimizers/base.py
imports: from trax import fastmath
  from trax.fastmath import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: clip_grads(grad_tree, max_norm) and returns: fastmath