code: def gather_tree_from_array(
      t: TensorLike, parent_ids: TensorLike, sequence_length: TensorLike
  ) -> tf.Tensor:
      """Calculates the full beams for a `TensorArray`.
  
      Args:
        t: A stacked `TensorArray` of size `max_time` that contains `Tensor`s of
          shape `[batch_size, beam_width, s]` or `[batch_size * beam_width, s]`
          where `s` is the depth shape.
        parent_ids: The parent ids of shape `[max_time, batch_size, beam_width]`.
        sequence_length: The sequence length of shape `[batch_size, beam_width]`.
  
      Returns:
        A `Tensor` which is a stacked `TensorArray` of the same size and type as
        `t` and where beams are sorted in each `Tensor` according to
        `parent_ids`.
      """
      max_time = parent_ids.shape[0] or tf.shape(parent_ids)[0]
      batch_size = parent_ids.shape[1] or tf.shape(parent_ids)[1]
      beam_width = parent_ids.shape[2] or tf.shape(parent_ids)[2]
  
      # Generate beam ids that will be reordered by gather_tree.
      beam_ids = tf.reshape(tf.range(beam_width), [1, 1, -1])
      beam_ids = tf.tile(beam_ids, [max_time, batch_size, 1])
  
      max_sequence_lengths = tf.cast(
          tf.reduce_max(sequence_length, axis=1), tf.int32
      )
      sorted_beam_ids = gather_tree(
          step_ids=beam_ids,
          parent_ids=parent_ids,
          max_sequence_lengths=max_sequence_lengths,
          end_token=beam_width + 1,
      )
  
      # For out of range steps, simply copy the same beam.
      in_bound_steps = tf.transpose(
          tf.sequence_mask(sequence_length, maxlen=max_time),
          perm=[2, 0, 1],
      )
      sorted_beam_ids = tf.where(
          in_bound_steps, x=sorted_beam_ids, y=beam_ids
      )
  
      # Gather from a tensor with collapsed additional dimensions.
      final_shape = tf.shape(t)
      gather_from = tf.reshape(
          t, [max_time, batch_size, beam_width, -1]
      )
      ordered = tf.gather(
          gather_from, sorted_beam_ids, axis=2, batch_dims=2
      )
      ordered = tf.reshape(ordered, final_shape)
  
      return ordered
  
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 5 for calling a global function with example usage: gather_tree_from_array(t, parent_ids, sequence_length) and returns: ordered