def fast_decode_tpu(encoder_output,
                    encoder_decoder_attention_bias,
                    symbols_to_logits_fn,
                    hparams,
                    decode_length,
                    vocab_size,
                    init_cache_fn=_init_transformer_cache,
                    beam_size=1,
                    top_beams=1,
                    alpha=1.0,
                    sos_id=0,
                    eos_id=beam_search.EOS_ID,
                    batch_size=None,
                    force_decode_length=False,
                    scope_prefix="body/",
                    use_top_k_with_unique=True,
                    sampling_temperature=0.0,
                    top_k=-1):
  """Given encoder output and a symbols to logits function, does fast decoding.

  Implements both greedy and beam search decoding for TPU, uses beam search iff
  beam_size > 1, otherwise beam search related arguments are ignored.

  Args:
    encoder_output: A tensor, output from encoder.
    encoder_decoder_attention_bias: A tensor, bias for use in encoder-decoder
      attention.
    symbols_to_logits_fn: Incremental decoding, function mapping triple `(ids,
      step, cache)` to symbol logits.
    hparams: Run hyperparameters.
    decode_length: An integer, how many additional timesteps to decode.
    vocab_size: Output vocabulary size.
    init_cache_fn: Function that returns the initial cache dict.
    beam_size: An integer, number of beams.
    top_beams: An integer, how many of the beams to return.
    alpha: A float that controls the length penalty. Larger the alpha, stronger
      the preference for longer translations.
    sos_id: Start-of-sequence symbol.
    eos_id: End-of-sequence symbol.
    batch_size: An integer, must be passed if there is no input.
    force_decode_length: A bool, whether to force the full decode length, or if
      False, stop when all beams hit eos_id.
    scope_prefix: str, prefix for decoder layer variable scopes.
    use_top_k_with_unique: bool, whether to use a fast (but decreased precision)
      top_k during beam search.
    sampling_temperature: scalar, temperature with which to sample.
    top_k: scalar, sample only top k.

  Returns:
    A dict of decoding results {
        "outputs": integer `Tensor` of decoded ids of shape
            [batch_size, <= decode_length] if top_beams == 1 or
            [batch_size, top_beams, <= decode_length] otherwise
        "scores": decoding log probs from the beam search,
            None if using greedy decoding (beam_size=1)
    }.

  Raises:
    NotImplementedError: If beam size > 1 with partial targets.
  """
  if encoder_output is not None:
    batch_size = common_layers.shape_list(encoder_output)[0]

  cache = init_cache_fn(None, hparams, batch_size, decode_length,
                        encoder_output, encoder_decoder_attention_bias,
                        scope_prefix)

  mlperf_log.transformer_print(
      key=mlperf_log.MODEL_HP_SEQ_BEAM_SEARCH,
      value={
          "vocab_size": vocab_size,
          "batch_size": batch_size,
          "beam_size": beam_size,
          "alpha": alpha,
          "max_decode_length": decode_length
      },
      hparams=hparams)
  if beam_size > 1:  # Beam Search
    initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
    decoded_ids, scores, _ = beam_search.beam_search(
        symbols_to_logits_fn,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        alpha,
        states=cache,
        eos_id=eos_id,
        stop_early=(top_beams == 1),
        use_tpu=True,
        use_top_k_with_unique=use_top_k_with_unique)

    if top_beams == 1:
      decoded_ids = decoded_ids[:, 0, 1:]
      scores = scores[:, 0]
    else:
      decoded_ids = decoded_ids[:, :top_beams, 1:]
      scores = scores[:, :top_beams]
  else:  # Greedy

    def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
      """One step of greedy decoding."""
      logits, cache = symbols_to_logits_fn(next_id, i, cache)
      log_probs = common_layers.log_prob_from_logits(logits)
      temperature = sampling_temperature
      if hparams.sampling_method == "random_per_example":
        next_id = common_layers.sample_temperature_per_example(
            logits, temperature, top_k)
      else:
        if hparams.sampling_method == "argmax":
          temperature = 0.0
        next_id = common_layers.sample_with_temperature(logits, temperature,
                                                        top_k)

      log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
                                  axis=1)
      log_prob += tf.gather_nd(
          log_probs, log_prob_indices) * (1 - tf.to_float(hit_eos))
      # Note(thangluong): we purposely update hit_eos after aggregating log_prob
      # There is a subtle detail here that we want to include log_probs up to
      # (and inclusive of) the first eos generated, but not subsequent tokens.
      hit_eos |= tf.equal(next_id, eos_id)

      next_id = tf.expand_dims(next_id, axis=1)
      decoded_ids = tf.transpose(decoded_ids)
      decoded_ids = inplace_ops.alias_inplace_update(
          decoded_ids, i, tf.squeeze(next_id, axis=1))
      decoded_ids = tf.transpose(decoded_ids)
      return i + 1, hit_eos, next_id, decoded_ids, cache, log_prob

    def is_not_finished(i, hit_eos, *_):
      finished = i >= decode_length
      if not force_decode_length:
        finished |= tf.reduce_all(hit_eos)
      return tf.logical_not(finished)

    decoded_ids = tf.zeros([batch_size, decode_length], dtype=tf.int64)
    hit_eos = tf.fill([batch_size], False)
    next_id = sos_id * tf.ones([batch_size, 1], dtype=tf.int64)
    initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)

    def compute_cache_shape_invariants(tensor):
      return tf.TensorShape(tensor.shape.as_list())

    _, _, _, decoded_ids, _, log_prob = tf.while_loop(
        is_not_finished,
        inner_loop, [
            tf.constant(0), hit_eos, next_id, decoded_ids, cache,
            initial_log_prob
        ],
        shape_invariants=[
            tf.TensorShape([]),
            tf.TensorShape([batch_size]),
            tf.TensorShape([batch_size, 1]),
            tf.TensorShape([batch_size, decode_length]),
            nest.map_structure(compute_cache_shape_invariants, cache),
            tf.TensorShape([batch_size]),
        ])
    scores = log_prob

  return {"outputs": decoded_ids, "scores": scores}
Given encoder output and a symbols to logits function, does fast decoding.

  Implements both greedy and beam search decoding for TPU, uses beam search iff
  beam_size > 1, otherwise beam search related arguments are ignored.

  Args:
    encoder_output: A tensor, output from encoder.
    encoder_decoder_attention_bias: A tensor, bias for use in encoder-decoder
      attention.
    symbols_to_logits_fn: Incremental decoding, function mapping triple `(ids,
      step, cache)` to symbol logits.
    hparams: Run hyperparameters.
    decode_length: An integer, how many additional timesteps to decode.
    vocab_size: Output vocabulary size.
    init_cache_fn: Function that returns the initial cache dict.
    beam_size: An integer, number of beams.
    top_beams: An integer, how many of the beams to return.
    alpha: A float that controls the length penalty. Larger the alpha, stronger
      the preference for longer translations.
    sos_id: Start-of-sequence symbol.
    eos_id: End-of-sequence symbol.
    batch_size: An integer, must be passed if there is no input.
    force_decode_length: A bool, whether to force the full decode length, or if
      False, stop when all beams hit eos_id.
    scope_prefix: str, prefix for decoder layer variable scopes.
    use_top_k_with_unique: bool, whether to use a fast (but decreased precision)
      top_k during beam search.
    sampling_temperature: scalar, temperature with which to sample.
    top_k: scalar, sample only top k.

  Returns:
    A dict of decoding results {
        "outputs": integer `Tensor` of decoded ids of shape
            [batch_size, <= decode_length] if top_beams == 1 or
            [batch_size, top_beams, <= decode_length] otherwise
        "scores": decoding log probs from the beam search,
            None if using greedy decoding (beam_size=1)
    }.

  Raises:
    NotImplementedError: If beam size > 1 with partial targets.
  
tensor2tensor/models/transformer.py