code: def _test_gather_tree_from_array(
      depth_ndims=0, merged_batch_beam=False
  ):
      array = np.array(
          [
              [[1, 2, 3], [4, 5, 6], [7, 8, 9], [0, 0, 0]],
              [[2, 3, 4], [5, 6, 7], [8, 9, 10], [11, 12, 0]],
          ]
      ).transpose([1, 0, 2])
      parent_ids = np.array(
          [
              [[0, 0, 0], [0, 1, 1], [2, 1, 2], [-1, -1, -1]],
              [[0, 0, 0], [1, 1, 0], [2, 0, 1], [0, 1, 0]],
          ]
      ).transpose([1, 0, 2])
      expected_array = np.array(
          [
              [[2, 2, 2], [6, 5, 6], [7, 8, 9], [0, 0, 0]],
              [[2, 3, 2], [7, 5, 7], [8, 9, 8], [11, 12, 0]],
          ]
      ).transpose([1, 0, 2])
      sequence_length = [[3, 3, 3], [4, 4, 3]]
  
      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
      )
  
      max_time = tf.shape(array)[0]
      batch_size = tf.shape(array)[1]
      beam_width = tf.shape(array)[2]
  
      def _tile_in_depth(tensor):
          # Generate higher rank tensors by concatenating tensor and
          # tensor + 1.
          for _ in range(depth_ndims):
              tensor = tf.stack([tensor, tensor + 1], -1)
          return tensor
  
      if merged_batch_beam:
          array = tf.reshape(array, [max_time, batch_size * beam_width])
          expected_array = tf.reshape(
              expected_array, [max_time, batch_size * beam_width]
          )
  
      if depth_ndims > 0:
          array = _tile_in_depth(array)
          expected_array = _tile_in_depth(expected_array)
  
      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 2 for calling a global function with example usage: _test_gather_tree_from_array(depth_ndims, merged_batch_beam) and returns: tensor