sub

mindspore/core/abstract/primitive_infer_map.cc

@@ -195,6 +195,7 @@ PrimitiveEvalImplMap &GetPrimitiveToBackendEvalImplMap() {
     {prim::kPrimDivNoNan, {InferImplDivNoNan, nullptr, true}},
     {prim::kPrimLinSpace, {InferImplLinSpace, nullptr, true}},
     {prim::kPrimAddN, {InferImplAddN, nullptr, true}},
+    {prim::kPrimSub, {InferImplSub, nullptr, false}},
 
     {prim::kPrimLess, {InferImplLess, nullptr, true}},
     {prim::kPrimStack, {InferImplStack, nullptr, true}},

mindspore/core/base/core_ops.h

@@ -45,6 +45,7 @@ constexpr auto kScalarTrunc = "ScalarTrunc";
 constexpr auto kScalarFloor = "ScalarFloor";
 constexpr auto kScalarUadd = "ScalarUadd";
 constexpr auto kScalarUsub = "ScalarUsub";
+constexpr auto kSub = "Sub";
 
 // Arrays
 constexpr auto kStack = "Stack";
@@ -371,7 +372,7 @@ inline const PrimitivePtr kPrimCentralization = std::make_shared<Primitive>("Cen
 inline const PrimitivePtr kPrimNeg = std::make_shared<Primitive>("Neg");
 inline const PrimitivePtr kPrimSin = std::make_shared<Primitive>("Sin");
 inline const PrimitivePtr kPrimCos = std::make_shared<Primitive>("Cos");
-inline const PrimitivePtr kPrimSub = std::make_shared<Primitive>("Sub");
+inline const PrimitivePtr kPrimSub = std::make_shared<Primitive>(kSub);
 inline const PrimitivePtr kPrimMul = std::make_shared<Primitive>("Mul");
 inline const PrimitivePtr kPrimDiv = std::make_shared<Primitive>("Div");
 inline const PrimitivePtr kPrimMod = std::make_shared<Primitive>("Mod");

mindspore/core/ops/op_utils.cc

@@ -25,14 +25,11 @@
 
 namespace mindspore {
 namespace ops {
-abstract::ShapePtr BroadCastInferShape(const std::string &op_name, const std::vector<AbstractBasePtr> &input_args) {
-  MS_LOG(INFO) << "Do infer shape for op " << op_name;
-  auto x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->GetShapeTrack())[kShape];
-  auto y_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[1]->GetShapeTrack())[kShape];
+std::vector<int64_t> CalBroadCastShape(std::vector<int64_t> x_shape, std::vector<int64_t> y_shape,
+                                       const std::string &op_name) {
   if (x_shape == y_shape) {
-    return std::make_shared<abstract::Shape>(x_shape);
+    return x_shape;
   }
-
   auto x_length = x_shape.size();
   auto y_length = y_shape.size();
   auto length = x_length < y_length ? x_length : y_length;
@@ -53,7 +50,29 @@ abstract::ShapePtr BroadCastInferShape(const std::string &op_name, const std::ve
       MS_EXCEPTION(ValueError) << "For op " << op_name << ", the two input can not broadcast";
     }
   }
-  return std::make_shared<abstract::Shape>(broadcast_shape);
+  return broadcast_shape;
+}
+abstract::ShapePtr BroadCastInferShape(const std::string &op_name, const std::vector<AbstractBasePtr> &input_args) {
+  MS_LOG(INFO) << "Do infer shape for op " << op_name;
+  CheckAndConvertUtils::CheckInteger("input numbers", input_args.size(), kGreaterEqual, 2, prim_name);
+  for (const auto &item : input_args) {
+    MS_EXCEPTION_IF_NULL(item);
+  }
+  auto x_shape_map = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->GetShapeTrack());
+  auto y_shape_map = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[1]->GetShapeTrack());
+  auto x_shape = x_shape_map[kShape];
+  auto y_shape = y_shape_map[kShape];
+  auto x_min_shape = x_shape_map[kMinShape];
+  auto x_max_shape = x_shape_map[kMaxShape];
+  auto y_min_shape = y_shape_map[kMinShape];
+  auto y_max_shape = y_shape_map[kMaxShape];
+  if (x_shape == y_shape) {
+    return std::make_shared<abstract::Shape>(x_shape, x_min_shape, x_max_shape);
+  }
+  auto broadcast_shape = CalBroadCastShape(x_shape, y_shape, op_name);
+  auto min_broadcast_shape = CalBroadCastShape(x_min_shape, y_min_shape, op_name);
+  auto max_broadcast_shape = CalBroadCastShape(x_max_shape, y_max_shape, op_name);
+  return std::make_shared<abstract::Shape>(broadcast_shape, min_broadcast_shape, max_broadcast_shape);
 }
 }  // namespace ops
 }  // namespace mindspore

mindspore/core/ops/sub.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020 Huawei Technologies Co., Ltd
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,12 +14,12 @@
  * limitations under the License.
  */
 
+#include "ops/sub.h"
 #include <map>
 #include <string>
 #include <vector>
 #include <memory>
 
-#include "ops/sub.h"
 #include "ops/op_utils.h"
 #include "utils/check_convert_utils.h"
 #include "abstract/primitive_infer_map.h"
@@ -30,6 +30,10 @@ namespace {
 abstract::ShapePtr InferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
   MS_EXCEPTION_IF_NULL(primitive);
   auto prim_name = primitive->name();
+  CheckAndConvertUtils::CheckInteger("input numbers", input_args.size(), kGreaterEqual, 2, prim_name);
+  for (const auto &item : input_args) {
+    MS_EXCEPTION_IF_NULL(item);
+  }
   return BroadCastInferShape(prim_name, input_args);
 }
 
@@ -37,6 +41,8 @@ TypePtr InferType(const PrimitivePtr &prim, const std::vector<AbstractBasePtr> &
   for (const auto &item : input_args) {
     MS_EXCEPTION_IF_NULL(item);
   }
+  auto op_name = prim->name();
+  CheckAndConvertUtils::CheckInteger("Sub infer", input_args.size(), kGreaterEqual, 2, op_name);
   std::map<std::string, TypePtr> types;
   types.emplace("x", input_args[0]->BuildType());
   types.emplace("y", input_args[1]->BuildType());

mindspore/core/ops/sub.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020 Huawei Technologies Co., Ltd
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,7 +26,7 @@
 
 namespace mindspore {
 namespace ops {
-constexpr auto kNameSub = "Sub";
+constexpr auto kNameSub = prim::kSub;
 class Sub : public PrimitiveC {
  public:
   Sub() : PrimitiveC(kNameSub) { InitIOName({"x", "y"}, {"output"}); }