import tensorflow as tf

    with tf.variable_scope(target_modality.name):
      log_info("Transforming body output with %s.top", target_modality.name)
      last_only = (
          target_modality.top_is_pointwise and
          self.hparams.mode == tf.estimator.ModeKeys.PREDICT and
          not self.hparams.force_full_predict)
      if not last_only:
        logits = target_modality.top(body_output, features["targets"])
      else:
        # Take body outputs for the last position only, and targets too.
        last_position_body_output = tf.expand_dims(
            body_output[:, -1, :, :], axis=[1])
        last_position_targets = tf.expand_dims(
            features["targets"][:, -1:, :, :], axis=[1])
        logits = target_modality.top(last_position_body_output,
                                     last_position_targets)
    return logits

  def top(self, body_output, features):
    if isinstance(body_output, dict):
      logits = {}
      for k, v in body_output.iteritems():
        logits[k] = self._top_single(v, features)
      return logits
