mindspore.mint.gather

mindspore.mint.gather(input, dim, index)

Gather data from a tensor by indices.

\[output[(i_0, i_1, ..., i_{dim}, i_{dim+1}, ..., i_n)] = input[(i_0, i_1, ..., index[(i_0, i_1, ..., i_{dim}, i_{dim+1}, ..., i_n)], i_{dim+1}, ..., i_n)]\]

Warning

On Ascend, the behavior is unpredictable in the following cases:

• the value of index is not in the range [-input.shape[dim], input.shape[dim]) in forward;

• the value of index is not in the range [0, input.shape[dim]) in backward.

Parameters

• input (Tensor) – The target tensor to gather values.

• dim (int) – the axis to index along, must be in range [-input.rank, input.rank).

• index (Tensor) –

  The index tensor, with int32 or int64 data type. An valid index should be:

  • index.rank == input.rank;

  • for axis != dim, index.shape[axis] <= input.shape[axis];

  • the value of index is in range [-input.shape[dim], input.shape[dim]).

Returns

Tensor, has the same type as input and the same shape as index.

Raises

• ValueError – If the shape of index is illegal.

• ValueError – If dim is not in [-input.rank, input.rank).

• ValueError – If the value of index is out of the valid range.

• TypeError – If the type of index is illegal.

Supported Platforms:

Ascend

Examples

>>> import mindspore
>>> import numpy as np
>>> from mindspore import Tensor, mint
>>> from mindspore.ops.function.array_func import gather_ext
>>> input = Tensor(np.array([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]), mindspore.float32)
>>> index = Tensor(np.array([[0, 0], [1, 1]]), mindspore.int32)
>>> output = gather_ext(input, 1, index)
>>> print(output)
[[-0.1 -0.1]
 [0.5   0.5]]