class_docstr: Adapts T2TModel for sequential_scheduled_sampling().
class_docstr_tok: ["Adapts","T2TModel","for","sequential_scheduled_sampling","(",")","."]
docstr: 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]
docstr_tok: ["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","]"]
code: 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
  
file_path: tensor2tensor/tensor2tensor/utils/scheduled_sampling.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.python.ops import inplace_ops
label_desc: Use image node_id 2 for calling the ScheduledSamplingAdapter obj's underlying member method code with example usage: obj.infer_fn(partial_targets) and returns: logits