mindspore.mint.index_select

mindspore.mint.index_select(input, dim, index)[source]

Generates a new tensor that accesses the values of input along the specified dim dimension using the indices specified in index. The new tensor has the same number of dimensions as input, with the size of the dim dimension equal to the length of index, and the sizes of all other dimensions unchanged from the original input tensor.

Note

The value of index must be in the range [0, input.shape[dim]). Negative or out-of-range values will lead to undefined behavior.

Parameters:

input (Tensor) – The input tensor.
dim (int) – The dimension to be indexed.
index (Tensor) – A 1-D tensor with the indices, whose dtype must be int32 or int64.

Returns:

Tensor, has the same dtype as the input tensor.

Raises:

TypeError – If input or index is not a Tensor.
TypeError – If dim is not an int.
ValueError – If the value of dim is out of the range of [-input.ndim, input.ndim - 1].
ValueError – If the dimension of index is not equal to 1.

Supported Platforms:: Ascend

Examples

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