!25643 fix zeros core dump when only has one input & batchmatmul input limit error

Merge pull request !25643 from Simson/op_infer
This commit is contained in:
i-robot 2021-10-30 01:48:14 +00:00 committed by Gitee
commit 289986fdfa
2 changed files with 4 additions and 2 deletions

View File

@ -36,8 +36,8 @@ abstract::ShapePtr BatchMatmulInferShape(const PrimitivePtr &primitive,
auto y_shape_map = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[1]->BuildShape());
auto x_shp = x_shape_map[kShape];
auto y_shp = y_shape_map[kShape];
constexpr size_t x_dim_limit = 2;
constexpr size_t y_dim_limit = 3;
constexpr size_t x_dim_limit = 3;
constexpr size_t y_dim_limit = 2;
if (x_shp.size() < x_dim_limit || y_shp.size() < y_dim_limit) {
MS_EXCEPTION(ValueError) << "For BatchMatMul, input x should be greater or equal to 3, input y should be greater "
"or equal to 2 while x size = "

View File

@ -57,6 +57,8 @@ AbstractBasePtr ZerosInfer(const abstract::AnalysisEnginePtr &, const PrimitiveP
ValuePtr ZerosInferValue(const PrimitivePtr &prim, const std::vector<AbstractBasePtr> &input_args) {
MS_EXCEPTION_IF_NULL(prim);
const int64_t input_num = 2;
CheckAndConvertUtils::CheckInputArgs(input_args, kGreaterEqual, input_num, prim->name());
auto abs = ZerosInfer(nullptr, prim, input_args);
// check
auto out_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(abs->BuildShape())[kShape];