!12199 getitem by int add restrict data rank in [0,8]

From: @yepei6
Reviewed-by: 
Signed-off-by:
This commit is contained in:
mindspore-ci-bot 2021-02-08 14:37:57 +08:00 committed by Gitee
commit 43eb2afd98
1 changed files with 7 additions and 0 deletions

View File

@ -161,6 +161,13 @@ def tensor_index_by_slice(data, slice_index):
def tensor_index_by_number(data, number):
"""Tensor getitem by a Number which may be integer/float/bool value"""
data_type = F.typeof(data)
if const_utils.judge_index_type(data_type, mstype.tensor_type):
data_shape = F.shape(data)
data_rank = len(data_shape)
min_data_rank, max_data_rank = 0, 8
const_utils.judge_data_rank(data_rank, min_data_rank, max_data_rank)
number_type = const_utils.check_number_index_type(number)
if number_type == const_utils.BOOL_:
return _tensor_index_by_bool(data, number)