code: def sequential_scheduled_sampling(infer_fn, mix_fn, loss_fn, targets):
      """Scheduled Sampling.
  
      Args:
        infer_fn: Function. Computes logits for all timesteps.
        mix_fn: Function. Mixes gold and sample tokens.
        loss_fn: Function. Computes loss between gold tokens and logits.
        targets: Tensor of shape [batch_size, seq_len]. Gold tokens.
  
      Returns:
        ss_tokens: Tensor of shape [batch_size, seq_len]. Scheduled sampling tokens.
        ss_logits: Tensor of shape [batch_size, seq_len, vocab_size]. Logits for
          next token when conditioning on ss_tokens.
        losses_dict: {str: scalar Tensor}. Losses to optimize.
      """
      targets_shape = common_layers.shape_list(targets)
      batch_size = targets_shape[0]
      seq_len = targets_shape[1]
  
      if not targets.shape.is_fully_defined():
          # TODO(duckworthd): When running on GPU, I get the following error. Solve
          # it to enable use on other devices.
          #
          #   Cannot use 'Identity_186' as input to
          #   'transformer/parallel_0_7/transformer/transformer/symbol_modality_16282_512/shared/convert_gradient_to_tensor_HBc3xYw22Mw'
          #   because 'Identity_186' is in a while loop.
  
          raise ValueError(
              "The following code only works on TPU. As targets.shape isn't fully "
              "defined, I am assuming you are using a different device."
          )
  
      def cond_fn(i, ss_tokens):
          """True if i < seq_len."""
          _ = ss_tokens
          return i < seq_len
  
      def body_fn(i, ss_tokens):
          """Constructs conditioning tokens for scheduled sampling."""
          # next_token_logits depends on timesteps 0...i-1.
          #
          # [batch_size, seq_len] -> [batch_size, seq_len, vocab_size]
          ss_tokens_logits = infer_fn(ss_tokens)
  
          # Same as 'next_token_logits = ss_tokens_logits[:, i, :]'.
          vocab_size = common_layers.shape_list(ss_tokens_logits)[2]
          next_token_logits = tf.slice(
              ss_tokens_logits,
              begin=[0, i, 0],
              size=[batch_size, 1, vocab_size],
          )
          next_token_logits = tf.squeeze(next_token_logits, axis=[1])
  
          # [batch_size, vocab_size] -> [batch_size]
          sampled_next_tokens = _sample_next_tokens(next_token_logits)
  
          # Same as 'gold_next_tokens = targets[:, i]'.
          gold_next_tokens = tf.slice(
              targets, begin=[0, i], size=[batch_size, 1]
          )
          gold_next_tokens = tf.squeeze(gold_next_tokens, axis=[1])
  
          next_tokens = mix_fn(gold_next_tokens, sampled_next_tokens)
          ss_tokens = _update_timestep(
              ss_tokens, timestep=i, values=next_tokens
          )
  
          return i + 1, tf.stop_gradient(ss_tokens)
  
      # tf.while_loop() over all timesteps. Generate scheduled sampling tokens.
      i = 0
      ss_tokens = tf.zeros([batch_size, seq_len], dtype=tf.int32)
      i, ss_tokens = tf.while_loop(cond_fn, body_fn, [i, ss_tokens])
  
      ss_logits = infer_fn(ss_tokens)
      return ss_tokens, ss_logits, loss_fn(targets, ss_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 a global function with example usage: sequential_scheduled_sampling(infer_fn, mix_fn, loss_fn, targets) and returns: ss_tokens,  ss_logits,  loss_fn, unknown, unknown,  tf