code: def test_gather_tree_from_array_complex_trajectory():
      # Max. time = 7, batch = 1, beam = 5.
      array = np.expand_dims(
          np.array(
              [
                  [[25, 12, 114, 89, 97]],
                  [[9, 91, 64, 11, 162]],
                  [[34, 34, 34, 34, 34]],
                  [[2, 4, 2, 2, 4]],
                  [[2, 3, 6, 2, 2]],
                  [[2, 2, 2, 3, 2]],
                  [[2, 2, 2, 2, 2]],
              ]
          ),
          -1,
      )
      parent_ids = np.array(
          [
              [[0, 0, 0, 0, 0]],
              [[0, 0, 0, 0, 0]],
              [[0, 1, 2, 3, 4]],
              [[0, 0, 1, 2, 1]],
              [[0, 1, 1, 2, 3]],
              [[0, 1, 3, 1, 2]],
              [[0, 1, 2, 3, 4]],
          ]
      )
      expected_array = np.expand_dims(
          np.array(
              [
                  [[25, 25, 25, 25, 25]],
                  [[9, 9, 91, 9, 9]],
                  [[34, 34, 34, 34, 34]],
                  [[2, 4, 2, 4, 4]],
                  [[2, 3, 6, 3, 6]],
                  [[2, 2, 2, 3, 2]],
                  [[2, 2, 2, 2, 2]],
              ]
          ),
          -1,
      )
      sequence_length = [[4, 6, 4, 7, 6]]
  
      array = tf.convert_to_tensor(array, dtype=tf.float32)
      parent_ids = tf.convert_to_tensor(parent_ids, dtype=tf.int32)
      expected_array = tf.convert_to_tensor(
          expected_array, dtype=tf.float32
      )
  
      sorted_array = beam_search_decoder.gather_tree_from_array(
          array, parent_ids, sequence_length
      )
  
      np.testing.assert_equal(
          expected_array.numpy(), sorted_array.numpy()
      )
  
imports: import numpy
  import pytest
  import tensorflow
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import beam_search_decoder, gather_tree
label_desc: Use image node_id 7 for calling a global function with example usage: test_gather_tree_from_array_complex_trajectory() without return types