code: def clip_by_global_norm(
      max_norm, use_psum=True
  ) -> GradientTransformation:
      """Clip updates using their global norm.
  
      References:
        [Pascanu et al, 2012](https://arxiv.org/abs/1211.5063)
  
      Args:
        max_norm: the maximum global norm for an update.
  
      Returns:
        An (init_fn, update_fn) tuple.
      """
  
      def init_fn(_):
          return ClipByGlobalNormState()
  
      def update_fn(updates, state, params=None):
          del params
          g_norm = global_norm(updates, use_psum=use_psum)
          trigger = g_norm < max_norm
          updates = jax.tree_map(
              lambda t: jnp.where(trigger, t, (t / g_norm) * max_norm),
              updates,
          )
          return updates, state
  
      return GradientTransformation(init_fn, update_fn)
  
imports: import jax
  import jax.numpy
  from jax.experimental.pjit import with_sharding_constraint
  from optax import AdditiveWeightDecayState, GradientTransformation, OptState
label_desc: Use image node_id 4 for calling a global function with example usage: clip_by_global_norm(max_norm, use_psum) and returns: GradientTransformation, ClipByGlobalNormState, updates,  state