code: def _make_weighted_reduce_sum(weights):
      """Builds a function to compute weighted sums."""
  
      def weighted_reduce_sum(x, axis=0):
          """Weighted sum over an axis of `x`."""
          # Extend the weights to broadcast over any event dimensions of `x`.
          # This assumes that `weights` and `x` have the same sample and batch
          # dimensions, e.g., that they come from the same `sample_and_log_prob` call.
          event_ndims = ps.rank(x) - ps.rank(weights)
          aligned_weights = tf.reshape(
              weights,
              ps.concat(
                  [
                      ps.shape(weights),
                      ps.ones([event_ndims], dtype=tf.int32),
                  ],
                  axis=0,
              ),
          )
          return tf.reduce_sum(
              aligned_weights * tf.cast(x, weights.dtype), axis=axis
          )
  
      return weighted_reduce_sum
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 1 for calling a global function with example usage: _make_weighted_reduce_sum(weights) and returns: weighted_reduce_sum, tf