Numerically stable version of log(reduce_sum(exp(x))).

  Unlike other reductions, the output has the same shape as the input.
  Note: with a minor change, we could allow multiple reduced dimensions.

  Args:
    x: a Tensor
    reduced_dim: a dimension in x
    extra_logit: an optional Tensor broadcastable to (x.shape - reduced_dim)
    name: an optional string
  Returns:
    a Tensor with the same shape and dtype as x.

def reduce_logsumexp(x, reduced_dim, extra_logit=None, name=None):
  """Numerically stable version of log(reduce_sum(exp(x))).

  Unlike other reductions, the output has the same shape as the input.
  Note: with a minor change, we could allow multiple reduced dimensions.

  Args:
    x: a Tensor
    reduced_dim: a dimension in x
    extra_logit: an optional Tensor broadcastable to (x.shape - reduced_dim)
    name: an optional string
  Returns:
    a Tensor with the same shape and dtype as x.
  """
  reduced_dim = convert_to_dimension(reduced_dim)
  with tf.variable_scope(name, default_name="reduce_logsumexp"):
    reduced_shape = x.shape - reduced_dim
    max_logit = reduce_max(stop_gradient(x), output_shape=reduced_shape)
    if extra_logit is not None:
      if isinstance(extra_logit, Tensor):
        extra_logit = stop_gradient(extra_logit)
      max_logit = maximum(max_logit, extra_logit)
    x -= max_logit
    exp_x = exp(x)
    sum_exp_x = reduce_sum(exp_x, output_shape=reduced_shape)
    if extra_logit is not None:
      sum_exp_x += exp(extra_logit - max_logit)
    return log(sum_exp_x) + max_logit