mindspore.ops.index_select

View Source On Gitee
mindspore.ops.index_select(input, axis, index)[source]

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

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.]]]