tfa.seq2seq.BeamSearchDecoderOutput

View source on GitHub

Outputs of a BeamSearchDecoder step.

Contains:

  • scores: The scores for this step, which are the log probabilities over the output vocabulary, possibly penalized by length and attention coverage. A float32 Tensor of shape [batch_size, beam_width, vocab_size].
  • predicted_ids: The token IDs predicted for this step. A int32 Tensor of shape [batch_size, beam_width].
  • parent_ids: The indices of the parent beam of each beam. A int32 Tensor of shape [batch_size, beam_width].

scores

predicted_ids

parent_ids