mindspore.mint.gather
- mindspore.mint.gather(input, dim, index)[源代码]
返回输入Tensor在指定 index 索引对应的元素组成的切片。
\[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)]\]警告
在Ascend后端,以下场景将导致不可预测的行为:
正向执行流程中,当 index 的取值不在范围 [-input.shape[dim], input.shape[dim]) 内;
反向执行流程中,当 index 的取值不在范围 [0, input.shape[dim]) 内。
- 参数:
input (Tensor) - 待索引切片取值的原始Tensor。
dim (int) - 指定要切片的维度索引。取值范围 [-input.rank, input.rank)。
index (Tensor) - 指定原始Tensor中要切片的索引。数据类型必须是int32或int64。需要同时满足以下条件:
index.rank == input.rank;
对于 axis != dim , index.shape[axis] <= input.shape[axis] ;
index 的取值在有效区间 [-input.shape[dim], input.shape[dim]) ;
- 返回:
Tensor,数据类型与 input 保持一致,shape与 index 保持一致。
- 异常:
ValueError - input 的shape取值非法。
ValueError - dim 取值不在有效范围 [-input.rank, input.rank)。
ValueError - index 的值不在有效范围。
TypeError - index 的数据类型非法。
- 支持平台:
Ascend
样例:
>>> import mindspore >>> import numpy as np >>> from mindspore import Tensor, mint >>> 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 = mint.gather(input, 1, index) >>> print(output) [[-0.1 -0.1] [0.5 0.5]]