gather_tree

paddle.fluid.layers.gather_tree(ids, parents)[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 (Variable) – 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 (Variable) – 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.

Return type

Variable

Examples

import paddle.fluid as fluid

ids = fluid.layers.data(name='ids',
                        shape=[5, 2, 2],
                        dtype='int64',
                        append_batch_size=False)
parents = fluid.layers.data(name='parents',
                            shape=[5, 2, 2],
                            dtype='int64',
                            append_batch_size=False)
final_sequences = fluid.layers.gather_tree(ids, parents)