使用mindspore.mint.gather函数计算出的结果错误

1 系统环境

硬件环境(Ascend/GPU/CPU): Ascend910
MindSpore版本: mindspore=2.4.0
执行模式(PyNative/ Graph): 不限
Python版本: Python=3.7
操作系统平台: linux

2 报错信息

2.1 问题描述

运行下述代码计算错误,程序正常运行,但给出错误结果。

2.2 脚本信息

import mindspore
from mindspore import mint, Tensor, ops, value_and_grad, context
import numpy as np

input_tensor = Tensor(np.array([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]), dtype = mindspore.float32)
index = Tensor(np.array([[0, 0], [1, 4]]), dtype = mindspore.int32)
output = mint.gather(input_tensor, 1, index)
print(output)

输出计算结果

[[-0.1 -0.1]
 [ 0.5  0.5]]

3 根因分析

gather操作返回输入Tensor在指定 index 索引对应的元素组成的切片。
input_tensor shape是(2,3)
index [[0, 0], [1, 4]]
其中4 取值不在范围 [0, input.shape[dim]) 内
api也有相关说明

警告  
在Ascend后端,以下场景将导致不可预测的行为:  
正向执行流程中,当 index 的取值不在范围 [-input.shape[dim], input.shape[dim]) 内;  
反向执行流程中,当 index 的取值不在范围 [0, input.shape[dim]) 内。

4 解决方案

修改index的值在有效范围内,例如修改成如下
index = Tensor(np.array([[0, 0], [1, 2]]), dtype = mindspore.int32)

import mindspore  
from mindspore import mint, Tensor, ops, value_and_grad, context  
import numpy as np  
    
input_tensor = Tensor(np.array([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]), dtype = mindspore.float32)  
index = Tensor(np.array([[0, 0], [1, 2]]), dtype = mindspore.int32)  
output = mint.gather(input_tensor, 1, index)  
print(output)