code: def test_gather_tree_one():
      # (max_time = 4, batch_size = 1, beams = 3)
      end_token = 10
      step_ids = _transpose_batch_time(
          [[[1, 2, 3], [4, 5, 6], [7, 8, 9], [-1, -1, -1]]]
      )
      parent_ids = _transpose_batch_time(
          [[[0, 0, 0], [0, 1, 1], [2, 1, 2], [-1, -1, -1]]]
      )
      max_sequence_lengths = [3]
      expected_result = _transpose_batch_time(
          [[[2, 2, 2], [6, 5, 6], [7, 8, 9], [10, 10, 10]]]
      )
      beams = gather_tree(
          step_ids=step_ids,
          parent_ids=parent_ids,
          max_sequence_lengths=max_sequence_lengths,
          end_token=end_token,
      )
      np.testing.assert_equal(expected_result, beams.numpy())
  
file_path: addons/tensorflow_addons/seq2seq/tests/beam_search_ops_test.py
imports: import itertools
  import numpy
  import pytest
  import tensorflow
  from tensorflow_addons.seq2seq import gather_tree
label_desc: Use image node_id 2 for calling a global function with example usage: test_gather_tree_one() without return types