code: def g_psum(x):
      return jax.lax.psum(x, "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 15 for calling a global function with example usage: g_psum(x) and returns: jax