code: def clip_gradients_vmap(g, l2_norm_clip):
      """Clips gradients in a way that is compatible with vectorized_map."""
      grads_flat = tf.nest.flatten(g)
      squared_l2_norms = [
          tf.reduce_sum(input_tensor=tf.square(g)) for g in grads_flat
      ]
      global_norm = tf.sqrt(tf.add_n(squared_l2_norms))
      div = tf.maximum(global_norm / l2_norm_clip, 1.0)
      clipped_flat = [g / div for g in grads_flat]
      clipped_grads = tf.nest.pack_sequence_as(g, clipped_flat)
      return clipped_grads
  
label_desc: Use image node_id 1 for calling a global function with example usage: clip_gradients_vmap(g, l2_norm_clip) and returns: clipped_grads