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 [0, input.shape[axis]). Negative or out-of-range values will lead to undefined behavior.
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.]]]