tfa.seq2seq.gather_tree

Calculates the full beams from the per-step ids and parent beam ids.

tfa.seq2seq.gather_tree(
    step_ids: tfa.types.TensorLike,
    parent_ids: tfa.types.TensorLike,
    max_sequence_lengths: tfa.types.TensorLike,
    end_token: tfa.types.Number
) -> tf.Tensor

For a given beam, past the time step containing the first decoded end_token all values are filled in with end_token.

Args
`step_ids`	The predicted token IDs. A `int32` `Tensor` of shape `[max_time, batch_size, beam_width]`.
`parent_ids`	The parent beam indices. A `int32` `Tensor` of shape `[max_time, batch_size, beam_width]`.
`max_sequence_lengths`	The maximum sequence length of each batch. A `int32` `Tensor` of shape `[batch_size]`.
`end_token`	The end token ID.

Returns
The reordered token IDs based on `parent_ids`.

Raises
`InvalidArgumentError`	if `parent_ids` contains an invalid index.

Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates.

Last updated 2023-05-25 UTC.

tfa.seq2seq.gather_treeStay organized with collections Save and categorize content based on your preferences.

Args

Returns

Raises

tfa.seq2seq.gather_tree