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