GatherTree

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:

for batch in range(BATCH_SIZE):
for beam in range(BEAM_WIDTH):
max_sequence_in_beam = min(MAX_TIME, max_seq_len[batch])
parent = parent_idx[max_sequence_in_beam - 1, batch, beam]
for level in reversed(range(max_sequence_in_beam - 1)):
final_idx[level, batch, beam] = step_idx[level, batch, parent]
parent = parent_idx[level, batch, parent]

Element data types for all input tensors should match each other.

Attributes: GatherTree has no attributes

Inputs

Outputs

Types

Example

<layer type="GatherTree" ...>
<input>
<port id="0">
<dim>100</dim>
<dim>1</dim>
<dim>10</dim>
</port>
<port id="1">
<dim>100</dim>
<dim>1</dim>
<dim>10</dim>
</port>
<port id="2">
<dim>1</dim>
</port>
<port id="3">
</port>
</input>
<output>
<port id="0">
<dim>100</dim>
<dim>1</dim>
<dim>10</dim>
</port>
</output>
</layer>