code: def global_norm(updates, use_psum=True):
      pre_sqrt = sum(
          [jnp.sum(jnp.square(x)) for x in jax.tree_leaves(updates)]
      )
      if use_psum:
          pre_sqrt = jax.lax.psum(pre_sqrt, "shard")
      return jnp.sqrt(pre_sqrt)
  
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 3 for calling a global function with example usage: global_norm(updates, use_psum) and returns: jnp