code: def rwb_psum(x, axis_name):
      """Applies a psum with reduce-without-broadcast (RWB) semantics.
  
      RWB semantics allow taking gradients w.r.t. unmapped variables of functions
      with psums in them.
  
      Args:
        x: a `Tensor` target for the psum.
        axis_name: A string axis name for the psum.
  
      Returns:
        A `Tensor` that is the result of applying a psum to an input `Tensor`.
      """
      return _rwb_psum_fwd(x, axis_name)[0]
  
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 14 for calling a global function with example usage: rwb_psum(x, axis_name) and returns: _rwb_psum_fwd