Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
0
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!34345 TensorScatterUpdate bool bug 

Merge pull request !34345 from ling/pr
This commit is contained in:
i-robot 2022-05-16 02:58:35 +00:00 committed by Gitee
parent 9118cf48ca c3d809b291
commit 9ae8279e40
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
2 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mindspore/core/ops/op_utils.h
View File

@ -29,6 +29,9 @@ namespace mindspore::ops {
const std::set<TypePtr> common_valid_types = {kInt8, kInt16, kInt32, kInt64, kUInt8, kUInt16,
kUInt32, kUInt64, kFloat16, kFloat32, kFloat64};
const std::set<TypePtr> common_valid_types_with_bool = {kInt8, kInt16, kInt32, kInt64, kUInt8, kUInt16,
kUInt32, kUInt64, kFloat16, kFloat32, kFloat64, kBool};
const std::set<TypePtr> common_valid_types_with_complex = {kInt8, kInt16, kInt32, kInt64, kUInt8,
kUInt16, kUInt32, kUInt64, kFloat16, kFloat32,
kFloat64, kComplex64, kComplex128};

3
mindspore/core/ops/tensor_scatter_arithmetic.cc
View File

@ -83,6 +83,9 @@ TypePtr TensorScatterArithmeticInferType(const PrimitivePtr &primitive,
std::map<std::string, TypePtr> type_dict;
type_dict.emplace("input_x", input_args[kInputIndex0]->BuildType());
type_dict.emplace("updates", input_args[kInputIndex2]->BuildType());
if (prim_name == prim::kPrimTensorScatterUpdate->name()) {
return CheckAndConvertUtils::CheckTensorTypeSame(type_dict, common_valid_types_with_bool, prim_name);
}
return CheckAndConvertUtils::CheckTensorTypeSame(type_dict, common_valid_types, prim_name);
}
} // namespace