code: def clip_grads(grad_tree, max_norm):
      """Clip gradients stored as a pytree of arrays to maximum norm `max_norm`."""
      # jax.example_libraries.optimizers.clip_grads implements this differently.
      # First, it uses clip_coef of max_norm / norm without the 1e-6.
      # Second, the jnp.where condition and argument order are reversed. This should
      # normally be a no-change but we do not know what changes in XLA are triggered
      # as a result of this and how that effects precision etc.
      norm = experimental_optimizers.l2_norm(grad_tree)
      clip_coef = max_norm / (norm + 1e-6)
      normalize = lambda g: jnp.where(clip_coef < 1.0, g * clip_coef, g)
      return jax.tree_util.tree_map(normalize, grad_tree)
  
imports: import copy
  import json
  import os
  from typing import Any, Dict, Optional, Set
  from absl import logging
  from flax import core
  from flax import optim
  from flax import traverse_util
  import jax
  from jax.example_libraries import optimizers
  import jax.numpy
  import numpy
  import PIL
  import PIL.ImageDraw
  import PIL.ImageFont
  from scenic.common_lib import image_utils
  from scenic.dataset_lib.coco_dataset import coco_eval
  from scenic.model_lib.base_models import box_utils
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import train_utils
  import scipy.special
  import tensorflow
label_desc: Use image node_id 8 for calling a global function with example usage: clip_grads(grad_tree, max_norm) and returns: jax