code: def reduce_logsumexp(x, reduced_dim, extra_logit=None, name=None):
    
    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
docs: 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.