parameters: (self, x, indices, batch_size)
return_statement: return nn.gather_nd(x, topk_coordinates)
docstring: r"""
          Gather from the tensor `x` using `indices`.
  
          Parameters:
              x(Variable): A tensor with shape `[batch_size, beam_size, ...]`.
              indices(Variable): A `int64` tensor with shape `[batch_size, beam_size]`,
                  representing the indices that we use to gather.
              batch_size(Variable): A tensor with shape `[1]`. Its data type should
                  be int32 or int64.
  
          Returns:
              Variable: A tensor with the same shape and data type as `x`, \
                  representing the gathered tensor.
docstring_summary: r"""
          Gather from the tensor `x` using `indices`.
docstring_tokens: r Gather from the tensor x using indices .
function: def _gather(self, x, indices, batch_size):
          r"""
          Gather from the tensor `x` using `indices`.
  
          Parameters:
              x(Variable): A tensor with shape `[batch_size, beam_size, ...]`.
              indices(Variable): A `int64` tensor with shape `[batch_size, beam_size]`,
                  representing the indices that we use to gather.
              batch_size(Variable): A tensor with shape `[1]`. Its data type should
                  be int32 or int64.
  
          Returns:
              Variable: A tensor with the same shape and data type as `x`, \
                  representing the gathered tensor.
          """
          check_type(x, 'x', (Variable), 'BeamSearchDecoder._gather')
          check_type(indices, 'indices', (Variable), 'BeamSearchDecoder._gather')
          check_type(batch_size, 'batch_size', (Variable),
                     'BeamSearchDecoder._gather')
          # TODO: compatibility of int32 and int64
          batch_size = tensor.cast(
              batch_size,
              indices.dtype) if batch_size.dtype != indices.dtype else batch_size
          batch_size.stop_gradient = True  # TODO: remove this
          batch_pos = paddle.tile(
              nn.unsqueeze(
                  tensor.range(
                      0, batch_size, 1, dtype=indices.dtype), [1]),
              [1, self.beam_size])
          topk_coordinates = nn.stack([batch_pos, indices], axis=2)
          topk_coordinates.stop_gradient = True
          return nn.gather_nd(x, topk_coordinates)
function_tokens: def _gather ( self , x , indices , batch_size ) : check_type ( x , 'x' , ( Variable ) , 'BeamSearchDecoder._gather' ) check_type ( indices , 'indices' , ( Variable ) , 'BeamSearchDecoder._gather' ) check_type ( batch_size , 'batch_size' , ( Variable ) , 'BeamSearchDecoder._gather' ) # TODO: compatibility of int32 and int64 batch_size = tensor . cast ( batch_size , indices . dtype ) if batch_size . dtype != indices . dtype else batch_size batch_size . stop_gradient = True # TODO: remove this batch_pos = paddle . tile ( nn . unsqueeze ( tensor . range ( 0 , batch_size , 1 , dtype = indices . dtype ) , [ 1 ] ) , [ 1 , self . beam_size ] ) topk_coordinates = nn . stack ( [ batch_pos , indices ] , axis = 2 ) topk_coordinates . stop_gradient = True return nn . gather_nd ( x , topk_coordinates )