mindspore.ops.index_select

mindspore.ops.index_select(input, axis, index)[source]

Select the input tensor according to the specified axis and index and return a new tensor.

Warning

After version 2.9.0, the parameter axis will be renamed to dim.

Note

The value of index must be in the range of [0, input.shape[axis]), the result is undefined out of range.
The returned tensor has the same number of dimensions as the input tensor. The axis dimension has the same size as the length of index , other dimensions have the same size as the input tensor.

Parameters

input (Tensor) – The input tensor.
axis (int) – The specified axis.
index (Tensor) – The specified indices, a 1-D tensor.

Returns

Tensor

Supported Platforms:: Ascend GPU CPU

Examples

>>> import mindspore
>>> input = mindspore.tensor(mindspore.ops.arange(0, 16).reshape(2, 2, 4), mindspore.float32)
>>> print(input)
[[[ 0.  1.  2.  3.]
  [ 4.  5.  6.  7.]]
 [[ 8.  9. 10. 11.]
  [12. 13. 14. 15.]]]
>>> index = mindspore.tensor([0,], mindspore.int32)
>>> y = mindspore.ops.index_select(input, 1, index)
>>> print(y)
[[[ 0.  1.  2.  3.]]
 [[ 8.  9. 10. 11.]]]