def test_gather_tree(self):
    # (max_time = 3, batch_size = 2, beam_width = 3)

    # create (batch_size, max_time, beam_width) matrix and transpose it
    predicted_ids = np.array(
        [[[1, 2, 3], [4, 5, 6], [7, 8, 9]], [[2, 3, 4], [5, 6, 7], [8, 9, 10]]],
        dtype=np.int32).transpose([1, 0, 2])
    parent_ids = np.array(
        [[[0, 0, 0], [0, 1, 1], [2, 1, 2]], [[0, 0, 0], [1, 2, 0], [2, 1, 1]]],
        dtype=np.int32).transpose([1, 0, 2])

    # sequence_lengths is shaped (batch_size = 3)
    max_sequence_lengths = [3, 3]

    expected_result = np.array([[[2, 2, 2], [6, 5, 6], [7, 8, 9]],
                                [[2, 4, 4], [7, 6, 6],
                                 [8, 9, 10]]]).transpose([1, 0, 2])

    res = beam_search_ops.gather_tree(
        predicted_ids,
        parent_ids,
        max_sequence_lengths=max_sequence_lengths,
        end_token=11)

    with self.cached_session() as sess:
      res_ = sess.run(res)

    self.assertAllEqual(expected_result, res_)

  def finalize(self, outputs, final_state, sequence_lengths):
    """Finalize and return the predicted_ids.

    Args:
      outputs: An instance of BeamSearchDecoderOutput.
      final_state: An instance of BeamSearchDecoderState. Passed through to the
        output.
      sequence_lengths: An `int64` tensor shaped `[batch_size, beam_width]`.
        The sequence lengths determined for each beam during decode.
        **NOTE** These are ignored; the updated sequence lengths are stored in
        `final_state.lengths`.

    Returns:
      outputs: An instance of `FinalBeamSearchDecoderOutput` where the
        predicted_ids are the result of calling _gather_tree.
      final_state: The same input instance of `BeamSearchDecoderState`.
    """
    del sequence_lengths
    # Get max_sequence_length across all beams for each batch.
    max_sequence_lengths = math_ops.to_int32(
        math_ops.reduce_max(final_state.lengths, axis=1))
    predicted_ids = beam_search_ops.gather_tree(
        outputs.predicted_ids,
        outputs.parent_ids,
        max_sequence_lengths=max_sequence_lengths,
        end_token=self._end_token)
    outputs = FinalBeamSearchDecoderOutput(
        beam_search_decoder_output=outputs, predicted_ids=predicted_ids)
    return outputs, final_state

def gather_tree_from_array(t, parent_ids, sequence_length):
  """Calculates the full beams for `TensorArray`s.

  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].value or array_ops.shape(parent_ids)[0]
  batch_size = parent_ids.shape[1].value or array_ops.shape(parent_ids)[1]
  beam_width = parent_ids.shape[2].value or array_ops.shape(parent_ids)[2]

  # Generate beam ids that will be reordered by gather_tree.
  beam_ids = array_ops.expand_dims(
      array_ops.expand_dims(math_ops.range(beam_width), 0), 0)
  beam_ids = array_ops.tile(beam_ids, [max_time, batch_size, 1])

  mask = array_ops.sequence_mask(
      sequence_length, maxlen=max_time, dtype=dtypes.int32)
  mask = array_ops.transpose(mask, perm=[2, 0, 1])

  # Use beam_width + 1 to mark the end of beam.
  masked_beam_ids = (beam_ids * mask) + (1 - mask) * (beam_width + 1)

  max_sequence_lengths = math_ops.to_int32(
      math_ops.reduce_max(sequence_length, axis=1))
  sorted_beam_ids = beam_search_ops.gather_tree(
      step_ids=masked_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.
  sorted_beam_ids = array_ops.where(
      math_ops.cast(mask, dtypes.bool), x=sorted_beam_ids, y=beam_ids)

  # Generate indices for gather_nd.
  time_ind = array_ops.tile(array_ops.reshape(
      math_ops.range(max_time), [-1, 1, 1]), [1, batch_size, beam_width])
  batch_ind = array_ops.tile(array_ops.reshape(
      math_ops.range(batch_size), [-1, 1, 1]), [1, max_time, beam_width])
  batch_ind = array_ops.transpose(batch_ind, perm=[1, 0, 2])
  indices = array_ops.stack([time_ind, batch_ind, sorted_beam_ids], -1)

  # Gather from a tensor with collapsed additional dimensions.
  gather_from = t
  final_shape = array_ops.shape(gather_from)
  gather_from = array_ops.reshape(
      gather_from, [max_time, batch_size, beam_width, -1])
  ordered = array_ops.gather_nd(gather_from, indices)
  ordered = array_ops.reshape(ordered, final_shape)

  return ordered

 def testGatherTreeOne(self):
   # (max_time = 4, batch_size = 1, beams = 3)
   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]]])
   sequence_length = [[3, 3, 3]]
   expected_result = _transpose_batch_time(
       [[[2, 2, 2], [6, 5, 6], [7, 8, 9], [-1, -1, -1]]])
   beams = beam_search_ops.gather_tree(
       step_ids=step_ids, parent_ids=parent_ids,
       sequence_length=sequence_length)
   with self.test_session(use_gpu=True):
     self.assertAllEqual(expected_result, beams.eval())

 def testBadParentValuesOnCPU(self):
   # (batch_size = 1, max_time = 4, beams = 3)
   # bad parent in beam 1 time 1
   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]]])
   sequence_length = [[3, 3, 3]]
   with ops.device("/cpu:0"):
     beams = beam_search_ops.gather_tree(
         step_ids=step_ids, parent_ids=parent_ids,
         sequence_length=sequence_length)
   with self.test_session():
     with self.assertRaisesOpError(
         r"parent id -1 at \(batch, time, beam\) == \(0, 0, 1\)"):
       _ = beams.eval()

 def testGatherTreeOne(self):
   # (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 = beam_search_ops.gather_tree(
       step_ids=step_ids,
       parent_ids=parent_ids,
       max_sequence_lengths=max_sequence_lengths,
       end_token=end_token)
   with self.cached_session(use_gpu=True):
     self.assertAllEqual(expected_result, self.evaluate(beams))

 def testBadParentValuesOnGPU(self):
   if not test.is_gpu_available():
     return
   # (max_time = 4, batch_size = 1, beams = 3)
   # bad parent in beam 1 time 1; appears as a negative index at time 0
   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]]])
   sequence_length = [3]
   expected_result = _transpose_batch_time(
       [[[2, -1, 2], [6, 5, 6], [7, 8, 9], [-1, -1, -1]]])
   with ops.device("/gpu:0"):
     beams = beam_search_ops.gather_tree(
         step_ids=step_ids, parent_ids=parent_ids,
         sequence_length=sequence_length)
   with self.test_session(use_gpu=True):
     self.assertAllEqual(expected_result, beams.eval())

 def testBadParentValuesOnCPU(self):
   # (batch_size = 1, max_time = 4, beams = 3)
   # bad parent in beam 1 time 1
   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]
   with ops.device("/cpu:0"):
     with self.assertRaisesOpError(
         r"parent id -1 at \(batch, time, beam\) == \(0, 0, 1\)"):
       beams = beam_search_ops.gather_tree(
           step_ids=step_ids,
           parent_ids=parent_ids,
           max_sequence_lengths=max_sequence_lengths,
           end_token=end_token)
       self.evaluate(beams)

  def finalize(self, outputs, final_state, sequence_lengths):
    """Finalize and return the predicted_ids.

    Args:
      outputs: An instance of BeamSearchDecoderOutput.
      final_state: An instance of BeamSearchDecoderState. Passed through to the
        output.
      sequence_lengths: An `int64` tensor shaped `[batch_size, beam_width]`.
        The sequence lengths determined for each beam during decode.

    Returns:
      outputs: An instance of FinalBeamSearchDecoderOutput where the
        predicted_ids are the result of calling _gather_tree.
      final_state: The same input instance of BeamSearchDecoderState.
    """
    predicted_ids = beam_search_ops.gather_tree(
        outputs.predicted_ids, outputs.parent_ids,
        sequence_length=sequence_lengths)
    outputs = FinalBeamSearchDecoderOutput(
        beam_search_decoder_output=outputs, predicted_ids=predicted_ids)
    return outputs, final_state

  def testGatherTreeBatch(self):
    batch_size = 10
    beam_width = 15
    max_time = 8
    max_sequence_lengths = [0, 1, 2, 4, 7, 8, 9, 10, 11, 0]
    end_token = 5

    with self.test_session(use_gpu=True):
      step_ids = np.random.randint(
          0, high=end_token + 1, size=(max_time, batch_size, beam_width))
      parent_ids = np.random.randint(
          0, high=beam_width - 1, size=(max_time, batch_size, beam_width))

      beams = beam_search_ops.gather_tree(
          step_ids=step_ids.astype(np.int32),
          parent_ids=parent_ids.astype(np.int32),
          max_sequence_lengths=max_sequence_lengths,
          end_token=end_token)

      self.assertEqual((max_time, batch_size, beam_width), beams.shape)
      beams_value = beams.eval()
      for b in range(batch_size):
        # Past max_sequence_lengths[b], we emit all end tokens.
        b_value = beams_value[max_sequence_lengths[b]:, b, :]
        self.assertAllClose(b_value, end_token * np.ones_like(b_value))
      for batch, beam in itertools.product(
          range(batch_size), range(beam_width)):
        v = np.squeeze(beams_value[:, batch, beam])
        if end_token in v:
          found_bad = np.where(v == -1)[0]
          self.assertEqual(0, len(found_bad))
          found = np.where(v == end_token)[0]
          found = found[0]  # First occurrence of end_token.
          # If an end_token is found, everything before it should be a
          # valid id and everything after it should be -1.
          if found > 0:
            self.assertAllEqual(
                v[:found - 1] >= 0, np.ones_like(v[:found - 1], dtype=bool))
          self.assertAllClose(v[found + 1:],
                              end_token * np.ones_like(v[found + 1:]))

 def testBadParentValuesOnGPU(self):
   # Only want to run this test on CUDA devices, as gather_tree is not
   # registered for SYCL devices.
   if not test.is_gpu_available(cuda_only=True):
     return
   # (max_time = 4, batch_size = 1, beams = 3)
   # bad parent in beam 1 time 1; appears as a negative index at time 0
   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, -1, 2], [6, 5, 6], [7, 8, 9], [-1, -1, -1]]])
   with ops.device("/device:GPU:0"):
     beams = beam_search_ops.gather_tree(
         step_ids=step_ids,
         parent_ids=parent_ids,
         max_sequence_lengths=max_sequence_lengths,
         end_token=10)
   with self.test_session(use_gpu=True):
     self.assertAllEqual(expected_result, beams.eval())

  def testGatherTreeBatch(self):
    # sequence_length is [batch_size, beam_width] = [4, 5]
    sequence_length = [[0] * 5, [1] * 5, [2] * 5, [3] * 5]

    with self.test_session(use_gpu=True):
      # (max_time = 4, batch_size = 4, beam_width = 5)
      step_ids = _transpose_batch_time(
          [[[3, 4, 0, 4, 0],
            [4, 2, 0, 3, 1],
            [1, 1, 3, 2, 2],
            [3, 1, 2, 3, 4]],
           [[3, 4, 0, 4, 0],
            [4, 2, 0, 3, 1],
            [1, 1, 3, 2, 2],
            [3, 1, 2, 3, 4]],
           [[1, 2, 3, 4, 2],
            [2, 1, 1, 3, 2],
            [3, 0, 1, 0, 0],
            [3, 4, 0, 2, 4]],
           [[0, 2, 2, 3, 1],
            [3, 2, 2, 2, 3],
            [3, 4, 3, 0, 3],
            [1, 2, 2, 2, 4]]])
      parent_ids = _transpose_batch_time(
          [[[4, 2, 4, 3, 4],
            [3, 4, 0, 2, 0],
            [3, 1, 3, 2, 2],
            [0, 2, 1, 4, 2]],
           [[4, 2, 4, 3, 4],
            [3, 4, 0, 2, 0],
            [3, 1, 3, 2, 2],
            [0, 2, 1, 4, 2]],
           [[3, 0, 0, 4, 0],
            [1, 2, 4, 2, 2],
            [4, 4, 0, 3, 0],
            [2, 4, 4, 3, 0]],
           [[3, 1, 4, 1, 3],
            [3, 2, 4, 0, 4],
            [1, 0, 1, 4, 2],
            [0, 3, 2, 0, 1]]])
      expected_beams = _transpose_batch_time(
          [[[-1, -1, -1, -1, -1],
            [-1, -1, -1, -1, -1],
            [-1, -1, -1, -1, -1],
            [-1, -1, -1, -1, -1]],
           [[3, 4, 0, 4, 0],
            [-1, -1, -1, -1, -1],
            [-1, -1, -1, -1, -1],
            [-1, -1, -1, -1, -1]],
           [[2, 3, 2, 3, 3],
            [2, 1, 1, 3, 2],
            [-1, -1, -1, -1, -1],
            [-1, -1, -1, -1, -1]],
           [[2, 3, 2, 1, 1],
            [2, 3, 2, 3, 2],
            [3, 4, 3, 0, 3],
            [-1, -1, -1, -1, -1]]])

      beams = beam_search_ops.gather_tree(
          step_ids=step_ids, parent_ids=parent_ids,
          sequence_length=sequence_length)
      self.assertAllEqual(expected_beams, beams.eval())