gather

paddle. gather ( x: Tensor, index: Tensor, axis: Tensor | int | None = None, name: str | None = None, out: Tensor | None = None ) Tensor [source]

This function has two functionalities, depending on the parameters passed:

  1. gather(Tensor input, int dim, Tensor index, Tensor out = None):

    PyTorch compatible gather, calls a non-broadcast paddle.take_along_axis. Check out take_along_axis and also [torch has more parameters] torch.scatter Note that sparse_grad param of PyTorch is currently not supported by Paddle, therefore do not pass this param (behavior is equivalent to sparse_grad = False). Also, dim allows for Tensor input, the same as PyTorch. However, when the first 3 params are all of Tensor types, there will be ambiguity between these two functionalities. Currently, original gather pass is more actively selected. Try avoiding using Tensor dim as input therefore.

  2. gather(Tensor x, Tensor index, int axis, str name = None, Tensor out = None):

    The original paddle.gather, see the following docs.

Output is obtained by gathering entries of axis of x indexed by index and concatenate them together.

Given:

x = [[1, 2],
     [3, 4],
     [5, 6]]

index = [1, 2]
axis=[0]

Then:

out = [[3, 4],
       [5, 6]]
Parameters
  • x (Tensor) – The source input tensor with rank>=1. Supported data type is int32, int64, float32, float64, complex64, complex128 and uint8 (only for CPU), float16 (only for GPU).

  • index (Tensor) – The index input tensor with rank=0 or rank=1. Data type is int32 or int64.

  • axis (Tensor|int|None, optional) – The axis of input to be gathered, it’s can be int or a Tensor with data type is int32 or int64. The default value is None, if None, the axis is 0.

  • name (str|None, optional) – The default value is None. Normally there is no need for user to set this property. For more information, please refer to api_guide_Name .

Returns

output (Tensor), If the index is a 1-D tensor, the output is a tensor with the same shape as x. If the index is a 0-D tensor, the output will reduce the dimension where the axis pointing.

Examples

>>> import paddle

>>> input = paddle.to_tensor([[1,2],[3,4],[5,6]])
>>> index = paddle.to_tensor([0,1])
>>> output = paddle.gather(input, index, axis=0)
>>> print(output)
Tensor(shape=[2, 2], dtype=int64, place=Place(cpu), stop_gradient=True,
[[1, 2],
 [3, 4]])