def maybe_scheduled_sampling(self, features, logits, losses):
    """Scheduled sampling.

    Performs forward inference again with "targets" feature replaced with values
    sampled from the model.

    This is the identity unless self.hparams.scheduled_sampling_prob > 0
    (default).

    **WARNING**: If hparams.scheduled_sampling_method == "parallel", this is
    not a faithful implementation of scheduled sampling. This implementation
    samples tokens for timestep t condtioned on gold tokens 1...t-1. A proper
    implementation must condition on a mix of gold and sampled tokens. Doing
    so is not efficient for models such like Transformer.

    Args:
      features: {str: Tensor}. Features sharded along batch dimension.
      logits: Tensor. Logits for each shard of data.
      losses: 0-D Tensor or (num: 0-D Tensor, denom: 0-D Tensor). Loss Tensor

    Returns:
      new_logits: Tensor.
      new_losses: {str: loss} where loss is one of (i) a 0-D Tensor or
        (ii) a (num: 0-D Tensor, denom: 0-D Tensor) pair to be used in a
        weighted average.
    """
    hparams = self.hparams
    problem_hparams = self._problem_hparams

    # Only do scheduled sampling if requested.
    if hparams.scheduled_sampling_prob == 0.0:
      return (logits, losses)

    # Only do scheduled sampling on language tasks.
    modality = problem_hparams.modality["targets"]
    if modality not in [
        modalities.ModalityType.SYMBOL,
        modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
        modalities.ModalityType.IMAGE
    ]:
      assert hparams.scheduled_sampling_prob == 0, (
          "Scheduled sampling only applies to ModalityType.(SYMBOL, "
          "SYMBOL_WEIGHTS_ALL, IMAGE). Found {modality}. Set "
          "hparams.scheduled_sampling_prob == 0.0.").format(modality=modality)
      return (logits, losses)

    # Only do scheduled sampling when training.
    is_training = (hparams.mode == tf_estimator.ModeKeys.TRAIN)
    if not is_training:
      tf.logging.info("Running in %s mode. Not using scheduled sampling.",
                      hparams.mode)
      return (logits, losses)

    # Pad vocabulary if vocab size must be evenly divisible by vocab_divisor.
    vocab_size = problem_hparams.vocab_size["targets"]
    assert vocab_size is not None
    assert hparams.vocab_divisor == 1

    # TODO(duckworthd): Move to scheduled_sampling.py.
    def sample(x):
      """Multinomial sampling from a n-dimensional tensor."""
      samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]), 1)
      reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
      return tf.to_int32(reshaped_samples)

    # TODO(duckworthd): Move to scheduled_sampling.py.
    def mix_gold_sampled(gold_targets,
                         sampled_targets,
                         mixin_prob,
                         i,
                         prev_new_targets):
      """Interleave sampled and gold tokens randomly."""
      # Resample each location iid.
      should_use_sampled_targets = tf.less(
          tf.random_uniform(common_layers.shape_list(sampled_targets)),
          mixin_prob)
      mixed_targets = tf.where(
          should_use_sampled_targets,
          sampled_targets,
          gold_targets)

      # Reuse sample tokens for earlier timesteps.
      new_targets = tf.where(
          is_later_timestep(gold_targets, i),
          mixed_targets,
          prev_new_targets)
      return new_targets

    # TODO(duckworthd): Move to scheduled_sampling.py.
    def is_later_timestep(x, pass_idx):
      """Constructs mask based on timestep."""
      assert x.shape.ndims == 4, x.shape
      x_shape = tf.shape(x)
      num_timesteps = x_shape[1]
      timesteps = tf.range(num_timesteps)
      timesteps = tf.reshape(timesteps, [1, num_timesteps, 1, 1])
      # The following is a bit untrue. For images, "num_timesteps" actually
      # represents image height, not time. We ignore that fact here.
      timesteps = tf.broadcast_to(timesteps, x_shape)
      return tf.greater_equal(timesteps, pass_idx)

    # TODO(duckworthd): Move to scheduled_sampling.py.
    def parallel_scheduled_sampling_pass(
        i, prev_new_targets, features, logits, mixin_prob):
      """Generate scheduled sampling results."""
      sampled_targets = sample(logits)
      new_targets = mix_gold_sampled(features["targets"],
                                     sampled_targets,
                                     mixin_prob,
                                     i,
                                     prev_new_targets)
      new_targets = tf.stop_gradient(new_targets)  # Treat new_targets as given.
      new_features = copy.copy(features)
      new_features["targets"] = new_targets
      with tf.variable_scope(tf.get_variable_scope(), reuse=True):
        # Compute bottom() for new_targets.
        #
        # TODO(duckworthd): Only apply bottom to 'new_targets'.
        new_transformed_features = self.bottom(new_features)

        # Compute body.
        with tf.variable_scope("body"):
          new_body_outputs, new_losses = self._normalize_body_output(
              self.body(new_transformed_features))
        assert "training" not in new_losses

        # Compute top.
        new_logits = self.top(new_body_outputs, new_features)

        # Compute loss. Use original features (== labels).
        if (hparams.mode != tf_estimator.ModeKeys.PREDICT and
            hparams.mode != "attack"):
          new_losses["training"] = self.loss(new_logits, features)
        else:
          new_losses["training"] = 0.0

      return new_targets, new_logits, new_losses

    tf.logging.info("Using scheduled sampling.")
    tf.logging.info("Warming scheduled sampling up with schedule: %s",
                    hparams.scheduled_sampling_warmup_schedule)
    assert hparams.scheduled_sampling_prob == 1.0, (
        "hparams.scheduled_sampling_prob must be 0 or 1.")

    if hparams.scheduled_sampling_method == "sequential":
      tf.logging.info("Using SEQUENTIAL scheduled sampling.")
      assert hparams.scheduled_sampling_num_passes == 1, (
          "hparams.scheduled_sampling_num_passes must equal 1 if "
          "doing sequential scheduled sampling.")
      return scheduled_sampling.sequential_scheduled_sampling_for_t2tmodel(
          self, features)
    elif hparams.scheduled_sampling_method == "parallel":
      tf.logging.info("Using PARALLEL scheduled sampling.")
      # TODO(duckworthd): Move this block to scheduled_sampling.py.

      # Gradually increase over a warmup period. Lower numbers mean more gold
      # tokens.
      mixin_prob = scheduled_sampling.inverse_decay_mix_prob(
          hparams.scheduled_sampling_warmup_schedule,
          hparams.scheduled_sampling_gold_mixin_prob,
          hparams.scheduled_sampling_warmup_steps)

      # Apply scheduled sampling over N passes. The logits from the (n-1)-th
      # pass will be mixed with gold tokens for conditioning in the n-th pass.
      assert hparams.scheduled_sampling_num_passes > 0, (
          "hparams.scheduled_sampling_num_passes must be > 0 if "
          "hparams.scheduled_sampling_prob > 0.0")
      new_logits = logits
      new_losses = losses
      prev_new_targets = features["targets"]
      for i in range(hparams.scheduled_sampling_num_passes):
        prev_new_targets, new_logits, new_losses = parallel_scheduled_sampling_pass(
            i, prev_new_targets, features, new_logits, mixin_prob)
      return new_logits, new_losses
    else:
      raise ValueError(
          "Unknown scheduled_sampling_method = %s" % (
              hparams.scheduled_sampling_method,))
Scheduled sampling.

    Performs forward inference again with "targets" feature replaced with values
    sampled from the model.

    This is the identity unless self.hparams.scheduled_sampling_prob > 0
    (default).

    **WARNING**: If hparams.scheduled_sampling_method == "parallel", this is
    not a faithful implementation of scheduled sampling. This implementation
    samples tokens for timestep t condtioned on gold tokens 1...t-1. A proper
    implementation must condition on a mix of gold and sampled tokens. Doing
    so is not efficient for models such like Transformer.

    Args:
      features: {str: Tensor}. Features sharded along batch dimension.
      logits: Tensor. Logits for each shard of data.
      losses: 0-D Tensor or (num: 0-D Tensor, denom: 0-D Tensor). Loss Tensor

    Returns:
      new_logits: Tensor.
      new_losses: {str: loss} where loss is one of (i) a 0-D Tensor or
        (ii) a (num: 0-D Tensor, denom: 0-D Tensor) pair to be used in a
        weighted average.
    
tensor2tensor/utils/t2t_model.py