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.

• 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)
#[[1. 2. 3. 4.]
# [5. 6. 7. 8.]
# [5. 6. 7. 8.]]
out_z2 = paddle.index_select(x=x, index=index, axis=1)
#[[ 1.  2.  2.]
# [ 5.  6.  6.]
# [ 9. 10. 10.]]