def METHOD_NAME():
    # (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())