Versioned name: GatherTree-1
Category: Beam search post-processing
Short description: Generates the complete beams from the ids per each step and the parent beam ids.
Detailed description
GatherTree operation implements the same algorithm as GatherTree operation in TensorFlow. Please see complete documentation here.
Pseudo code:
Element data types for all input tensors should match each other.
Attributes: GatherTree has no attributes
Inputs
step_ids
– a tensor of shape [MAX_TIME, BATCH_SIZE, BEAM_WIDTH]
of type T
with indices from per each step. Required.parent_idx
– a tensor of shape [MAX_TIME, BATCH_SIZE, BEAM_WIDTH]
of type T
with parent beam indices. Required.max_seq_len
– a tensor of shape [BATCH_SIZE]
of type T
with maximum lengths for each sequence in the batch. Required.end_token
– a scalar tensor of type T
with value of the end marker in a sequence. Required.Outputs
final_idx
– a tensor of shape [MAX_TIME, BATCH_SIZE, BEAM_WIDTH]
of type T
.Types
float32
or int32
; float32
should have integer values only.Example