gather_tree

paddle.nn.functional. gather_tree ( ids: Tensor, parents: Tensor ) → Tensor [source]

To be used after beam search. After beam search, we get selected ids at each time step and the corresponding parents in the search tree. Both ids and parents have the layout [max_time, batch_size, beam_size]. Then gather_tree is used to backtrace from the last time step and generate the full sequences by collecting selected ids.

Here is an example:

Given:
    ids = [[[2 2]
            [6 1]]
           [[3 9]
            [6 1]]
           [[0 1]
            [9 0]]]
    parents = [[[0 0]
                [1 1]]
               [[1 0]
                [1 0]]
               [[0 0]
                [0 1]]]

Then:
    gather_tree(ids, parents)
             = [[[2 2]
                 [1 6]]
                [[3 3]
                 [6 1]]
                [[0 1]
                 [9 0]]]

Parameters

• ids (Tensor) – A Tensor with shape [length, batch_size, beam_size] and data type int32 or int64. It contains the selected ids of all time steps.

• parents (Tensor) – A Tensor with the same shape and data type as ids, It contains the parents corresponding to selected ids when searching among beams.

Returns

A Tensor with the same shape and data type as ids. It contains the full sequences. The sequences are collected from ids by backtracing according to parents.

Examples

>>> import paddle

>>> ids = paddle.to_tensor([[[2, 2], [6, 1]], [[3, 9], [6, 1]], [[0, 1], [9, 0]]])

>>> parents = paddle.to_tensor([[[0, 0], [1, 1]], [[1, 0], [1, 0]], [[0, 0], [0, 1]]])

>>> final_sequences = paddle.nn.functional.gather_tree(ids, parents)
>>> [[[2, 2], [1, 6]], [[3, 3], [6, 1]], [[0, 1], [9, 0]]]
>>> final_sequences = paddle.nn.functional.gather_tree(ids, parents)
>>> print(final_sequences)
Tensor(shape=[3, 2, 2], dtype=int64, place=Place(cpu), stop_gradient=True,
[[[2, 2],
  [1, 6]],
 [[3, 3],
  [6, 1]],
 [[0, 1],
  [9, 0]]])