index_select

paddle. index_select ( x, index, axis=0, name=None ) [source]

Returns a new tensor which indexes the input tensor along dimension axis using the entries in index which is a Tensor. The returned tensor has the same number of dimensions as the original x tensor. The dim-th dimension has the same size as the length of index; other dimensions have the same size as in the x tensor.

Parameters
  • x (Tensor) – The input Tensor to be operated. The data of x can be one of float16, float32, float64, int32, int64, complex64 and complex128.

  • index (Tensor) – The 1-D Tensor containing the indices to index. The data type of index must be int32 or int64.

  • axis (int, optional) – The dimension in which we index. Default: if None, the axis is 0.

  • name (str, optional) – For details, please refer to Name. Generally, no setting is required. Default: None.

Returns

A Tensor with same data type as x.

Return type

Tensor

Examples

>>> import paddle

>>> x = paddle.to_tensor([[1.0, 2.0, 3.0, 4.0],
...                       [5.0, 6.0, 7.0, 8.0],
...                       [9.0, 10.0, 11.0, 12.0]])
>>> index = paddle.to_tensor([0, 1, 1], dtype='int32')
>>> out_z1 = paddle.index_select(x=x, index=index)
>>> print(out_z1.numpy())
[[1. 2. 3. 4.]
 [5. 6. 7. 8.]
 [5. 6. 7. 8.]]
>>> out_z2 = paddle.index_select(x=x, index=index, axis=1)
>>> print(out_z2.numpy())
[[ 1.  2.  2.]
 [ 5.  6.  6.]
 [ 9. 10. 10.]]