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)
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.
  
tensor2tensor/utils/scheduled_sampling.py