code: def reduce_logsumexp(
      x, axis=None, named_axis=None, allow_all_gather=False, **kwargs
  ):
      """`logsumexp` wrapper."""
      xmax = reduce_max(
          tf.stop_gradient(x),
          axis=axis,
          named_axis=named_axis,
          keepdims=True,
          allow_all_gather=allow_all_gather,
      )
      xmax = tf.where(
          tf.math.is_finite(xmax), xmax, tf.zeros_like(xmax)
      )
      result = tf.math.log(
          reduce_sum(
              tf.exp(x - xmax),
              axis=axis,
              named_axis=named_axis,
              **kwargs,
          )
      )
      return tf.reshape(xmax, ps.shape(result)) + result
  
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 8 for calling a global function with example usage: reduce_logsumexp(x, axis, named_axis, allow_all_gather) and returns: unknown