code: def f_pmean_bwd(_, g):
      return (jax.lax.pmean(g, "shard"),)
  
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 14 for calling a global function with example usage: f_pmean_bwd(_, g) and returns: jax