openvino.runtime.opset15.gather_tree#

openvino.runtime.opset15.gather_tree(step_ids: Node | int | float | ndarray, parent_idx: Node | int | float | ndarray, max_seq_len: Node | int | float | ndarray, end_token: Node | int | float | ndarray, name: str | None = None) → Node#

Perform GatherTree operation.

Parameters:

  • step_ids – The tensor with indices from per each step.

  • parent_idx – The tensor with with parent beam indices.

  • max_seq_len – The tensor with maximum lengths for each sequence in the batch.

  • end_token – The scalar tensor with value of the end marker in a sequence.

  • name – Optional name for output node.

Returns:

The new node performing a GatherTree operation.

The GatherTree node generates the complete beams from the indices per each step and the parent beam indices. GatherTree uses the following logic:

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]