log(softmax(x)).

  Args:
    x: a Tensor whose shape contains vocab_dim
    reduced_dim: a Dimension
    extra_logit: an optional Tensor broadcastable to (x.shape - reduced_dim)
    name: an optional string

  Returns:
    a Tensor with the same shape as x

def log_softmax(x, reduced_dim, extra_logit=None, name=None):
  """log(softmax(x)).

  Args:
    x: a Tensor whose shape contains vocab_dim
    reduced_dim: a Dimension
    extra_logit: an optional Tensor broadcastable to (x.shape - reduced_dim)
    name: an optional string

  Returns:
    a Tensor with the same shape as x
  """
  return x - reduce_logsumexp(
      x, reduced_dim, extra_logit=extra_logit, name=name)