code: def gather_tree(
      step_ids: TensorLike,
      parent_ids: TensorLike,
      max_sequence_lengths: TensorLike,
      end_token: Number,
  ) -> tf.Tensor:
      """Calculates the full beams from the per-step ids and parent beam ids.
  
      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.
      """
      if not options.is_custom_kernel_disabled():
          try:
              return _beam_search_so.ops.addons_gather_tree(
                  step_ids, parent_ids, max_sequence_lengths, end_token
              )
          except tf.errors.NotFoundError:
              options.warn_fallback("gather_tree")
  
      step_ids = tf.convert_to_tensor(step_ids, dtype=tf.int32)
      parent_ids = tf.convert_to_tensor(parent_ids, dtype=tf.int32)
      max_sequence_lengths = tf.convert_to_tensor(
          max_sequence_lengths, dtype=tf.int32
      )
      end_token = tf.convert_to_tensor(end_token, dtype=tf.int32)
      return _gather_tree(
          step_ids, parent_ids, max_sequence_lengths, end_token
      )
  
imports: import collections
  import numpy
  import tensorflow
  from tensorflow_addons import options
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import decoder
  from tensorflow_addons.utils import keras_utils
  from tensorflow_addons.utils.resource_loader import LazySO
  from tensorflow_addons.utils.types import FloatTensorLike, TensorLike, Number
  from typeguard import typechecked
  from typing import Callable, Optional
label_desc: Use image node_id 4 for calling a global function with example usage: gather_tree(step_ids, parent_ids, max_sequence_lengths, end_token) and returns: _gather_tree, _beam_search_so