'Generate logits.
Args:
body_output: A Tensor with shape [batch, p0, p1, body_input_depth]
Returns:
logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].'

def top(self, body_output, _):

    if self._model_hparams.shared_embedding_and_softmax_weights:
        scope_name = 'shared'
        reuse = True
    else:
        scope_name = 'softmax'
        reuse = False
    with tf.variable_scope(scope_name, reuse=reuse):
        var = self._get_weights()
        shape = tf.shape(body_output)[:(-1)]
        body_output = tf.reshape(body_output, [(-1), self._body_input_depth])
        logits = tf.matmul(body_output, var, transpose_b=True)
        logits = tf.reshape(logits, tf.concat([shape, [self._vocab_size]], 0))
        return tf.expand_dims(logits, 3)
