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

fix shape operator

This commit is contained in:
l00591931 2021-04-22 09:33:32 +08:00
parent ca8ab21233
commit c4537b6343
2 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/core/ops/shape.cc
View File

@ -36,6 +36,8 @@ AbstractBasePtr ShapeInfer(const abstract::AnalysisEnginePtr &, const PrimitiveP
auto shape_map = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape()); auto shape_map = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape());
auto in_shape = shape_map[kShape]; auto in_shape = shape_map[kShape];
// infer type // infer type
std::set<TypePtr> valid_params_types = {kTensorType};
CheckAndConvertUtils::CheckSubClass("shape type", input_args[0]->BuildType(), valid_params_types, op_name);
AbstractBasePtrList abs_list; AbstractBasePtrList abs_list;
std::transform(in_shape.begin(), in_shape.end(), std::back_inserter(abs_list), std::transform(in_shape.begin(), in_shape.end(), std::back_inserter(abs_list),
[](int64_t item) -> std::shared_ptr<abstract::AbstractScalar> { [](int64_t item) -> std::shared_ptr<abstract::AbstractScalar> {

2
mindspore/core/ops/softmax.cc
View File

@ -44,6 +44,7 @@ void Softmax::Init(const int64_t axis) {
this->set_axis(axis_vec); this->set_axis(axis_vec);
} }
namespace {
abstract::ShapePtr SoftMaxInferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) { abstract::ShapePtr SoftMaxInferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
MS_EXCEPTION_IF_NULL(primitive); MS_EXCEPTION_IF_NULL(primitive);
auto op_name = primitive->name(); auto op_name = primitive->name();
@ -74,6 +75,7 @@ TypePtr SoftMaxInferType(const PrimitivePtr &prim, const std::vector<AbstractBas
const std::set<TypePtr> valid_types = {kFloat16, kFloat32, kFloat64}; const std::set<TypePtr> valid_types = {kFloat16, kFloat32, kFloat64};
return CheckAndConvertUtils::CheckTensorTypeValid("x", input_args[0]->BuildType(), valid_types, prim->name()); return CheckAndConvertUtils::CheckTensorTypeValid("x", input_args[0]->BuildType(), valid_types, prim->name());
} }
} // namespace
AbstractBasePtr SoftmaxInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, AbstractBasePtr SoftmaxInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) { const std::vector<AbstractBasePtr> &input_args) {