code: def _rwb_psum_fwd(x, axis_name):
      if JAX_MODE:
          axis_name = canonicalize_named_axis(axis_name)
          out = lax.psum(x, axis_name)
      else:
          ctx = tf.distribute.get_replica_context()
          out = ctx.all_reduce("sum", x)
      return out, None
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow.python.util import nest
label_desc: Use image node_id 11 for calling a global function with example usage: _rwb_psum_fwd(x, axis_name) and returns: out,  None