Class ov::op::v1::GatherTree

class GatherTree : public ov::op::Op

Generates the complete beams from the ids per each step and the parent beam ids.

Public Functions

GatherTree(const Output<Node> &step_ids, const Output<Node> &parent_idx, const Output<Node> &max_seq_len, const Output<Node> &end_token)
Parameters
  • step_idsTensor of shape [MAX_TIME, BATCH_SIZE, BEAM_WIDTH] with indices from per each step

  • parent_idxTensor of shape [MAX_TIME, BATCH_SIZE, BEAM_WIDTH] with parent beam indices

  • max_seq_lenTensor of shape [BATCH_SIZE] with maximum lengths for each sequence in the batch

  • end_tokenTensor of shape [MAX_TIME, BATCH_SIZE, BEAM_WIDTH]

virtual void validate_and_infer_types() override

Verifies that attributes and inputs are consistent and computes output shapes and element types. Must be implemented by concrete child classes so that it can be run any number of times.

Throws if the node is invalid.