docs: def predict(self, encoder_outputs, encoder_decoder_attention_bias):
      """Return predicted sequence."""
      batch_size = tf.shape(encoder_outputs)[0]
      input_length = tf.shape(encoder_outputs)[1]
      max_decode_length = input_length + self.params.extra_decode_length
  
      symbols_to_logits_fn = self._get_symbols_to_logits_fn(max_decode_length)
  
      # Create initial set of IDs that will be passed into symbols_to_logits_fn.
      initial_ids = tf.zeros([batch_size], dtype=tf.int32)
  
      # Create cache storing decoder attention values for each layer.
      cache = {
          "layer_%d" % layer: {
              "k": tf.zeros([batch_size, 0, self.params.hidden_size]),
              "v": tf.zeros([batch_size, 0, self.params.hidden_size]),
          } for layer in range(self.params.num_hidden_layers)}
  
      # Add encoder output and attention bias to the cache.
      cache["encoder_outputs"] = encoder_outputs
      cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias
  
      # Use beam search to find the top beam_size sequences and scores.
      mlperf_log.transformer_print(
        key=mlperf_log.MODEL_HP_SEQ_BEAM_SEARCH,
        value={
          "vocab_size": self.params.vocab_size,
          "beam_size": self.params.beam_size,
          "alpha": self.params.alpha,
          "extra_decode_length": self.params.extra_decode_length})
      decoded_ids, scores = beam_search.sequence_beam_search(
          symbols_to_logits_fn=symbols_to_logits_fn,
          initial_ids=initial_ids,
          initial_cache=cache,
          vocab_size=self.params.vocab_size,
          beam_size=self.params.beam_size,
          alpha=self.params.alpha,
          max_decode_length=max_decode_length,
          eos_id=EOS_ID)
  
      # Get the top sequence for each batch element
      top_decoded_ids = decoded_ids[:, 0, 1:]
      top_scores = scores[:, 0]
  
      return {"outputs": top_decoded_ids, "scores": top_scores}