code: def f_psum_fwd(x):
      return f_psum(x), None
  
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 10 for calling a global function with example usage: f_psum_fwd(x) and returns: f_psum,  None