'Transform inputs from model space to target space.
Perform the Xception "Exit flow", consisting of a single residual block and
two separable convolutional upscalings followed by global spatial average
pooling.
Args:
body_output: A Tensor with shape [batch, ?, ?, body_output_size].
Returns:
a Tensors, each with shape [batch_size, ?, ?, vocab_size]'

def top(self, body_output, _):

    with tf.variable_scope(self.name):
        x = body_output
        if self._is_2d:
            length_float = tf.to_float(tf.shape(x)[1])
            spatial_dim_float = tf.sqrt(length_float)
            spatial_dim = tf.to_int32(spatial_dim_float)
            x = tf.reshape(x, [(-1), spatial_dim, spatial_dim, self._body_input_depth])
        x = common_layers.conv_block_downsample(x, self._kernel, self._strides, self._padding)
        x = tf.nn.relu(x)
        x = tf.reduce_mean(x, axis=[1, 2], keep_dims=True)
        res = common_layers.conv(x, self._vocab_size, (1, 1))
        return tf.expand_dims(res, 3)
