def infer_fn(self, partial_targets):
    """Computes logits for all timesteps.

    Args:
      partial_targets: [batch_size, seq_len]. Targets to condition on.

    Returns:
      next_token_logits: [batch_size, seq_len, vocab_size]
    """
    batch_size, seq_len = common_layers.shape_list(partial_targets)
    partial_targets = tf.reshape(partial_targets, [batch_size, seq_len, 1, 1])
    features = copy.copy(self._features)
    features["targets"] = partial_targets

    with tf.variable_scope(tf.get_variable_scope(), reuse=True):
      transformed_features = self._t2tmodel.bottom(features)

      with tf.variable_scope("body"):
        body_outputs, losses = self._t2tmodel._normalize_body_output(  # pylint: disable=protected-access
            self._t2tmodel.body(transformed_features))
        assert losses == {"extra": 0.0}, (
            "Auxiliary losses are not propagated in this code. %s"
            % (losses,))

      logits = self._t2tmodel.top(body_outputs, features)

    vocab_size = self._t2tmodel.problem_hparams.vocab_size["targets"]
    logits = tf.reshape(logits, [batch_size, seq_len, vocab_size])
    return logits
Computes logits for all timesteps.

    Args:
      partial_targets: [batch_size, seq_len]. Targets to condition on.

    Returns:
      next_token_logits: [batch_size, seq_len, vocab_size]
    
tensor2tensor/utils/scheduled_sampling.py