Remove "static inputs" for reduction ops

Author: Abhishek Kulkarni

ngraph_bridge/ngraph_builder.cc

@@ -2400,45 +2400,8 @@ static Status TranslateNonMaxSuppressionV4Op(
   return Status::OK();
 }
 
-static Status TranslateReduceOp(
-    const Node* op, const std::vector<const Tensor*>& static_input_map,
-    Builder::OpMap& ng_op_map,
-    std::function<std::shared_ptr<ng::Node>(
-        std::shared_ptr<ng::Node>, std::shared_ptr<ng::Node>, const bool)>
-        create_ng_node) {
-  shared_ptr<ng::Node> ng_input;
-  TF_RETURN_IF_ERROR(GetInputNode(ng_op_map, op, 0, &ng_input));
-  bool tf_keep_dims;
-  if (GetNodeAttr(op->attrs(), "keep_dims", &tf_keep_dims) != Status::OK()) {
-    tf_keep_dims = false;
-  }
-
-  std::vector<int64> axes;
-  TF_RETURN_IF_ERROR(GetStaticInputVector(op, 1, static_input_map, &axes));
-
-  ng::Shape input_shape = ng_input->get_shape();
-  size_t input_rank = input_shape.size();
-
-  TF_RETURN_IF_ERROR(CheckAxisDimInRange(axes, input_rank));
-
-  std::vector<size_t> ng_reduction_axes_vect(axes.size());
-  std::transform(
-      axes.begin(), axes.end(), ng_reduction_axes_vect.begin(),
-      [input_rank](int idx) { return idx + (idx < 0 ? (int)input_rank : 0); });
-  auto ng_reduction_axes = ConstructNgNode<ng::opset3::Constant>(
-      op->name(), ng::element::i64, ng::Shape{ng_reduction_axes_vect.size()},
-      ng_reduction_axes_vect);
-
-  std::shared_ptr<ng::Node> ng_node =
-      create_ng_node(ng_input, ng_reduction_axes, tf_keep_dims);
-  Builder::SetTracingInfo(op->name(), ng_node);
-
-  SaveNgOp(ng_op_map, op->name(), ng_node);
-  return Status::OK();
-}
-
 template <typename T>
-static Status TranslateDirectReduceOp(
+static Status TranslateReduceOp(
     const Node* op, const std::vector<const Tensor*>& static_input_map,
     Builder::OpMap& ng_op_map) {
   // ensure its either an arithmetic or a logical reduction
@@ -2448,13 +2411,19 @@ static Status TranslateDirectReduceOp(
         "Expected node to be either a valid logical or arithmetic reduction "
         "type");
   }
-  return TranslateReduceOp(
-      op, static_input_map, ng_op_map,
-      [&op](std::shared_ptr<ng::Node> ng_input,
-            std::shared_ptr<ng::Node> ng_reduction_axes, const bool keep_dims) {
-        return ConstructNgNode<T>(op->name(), ng_input, ng_reduction_axes,
-                                  keep_dims);
-      });
+
+  shared_ptr<ng::Node> ng_input, ng_reduction_indices;
+  TF_RETURN_IF_ERROR(
+      GetInputNodes(ng_op_map, op, &ng_input, &ng_reduction_indices));
+  bool keep_dims;
+  if (GetNodeAttr(op->attrs(), "keep_dims", &keep_dims) != Status::OK()) {
+    keep_dims = false;
+  }
+
+  std::shared_ptr<ng::Node> ng_node =
+      ConstructNgNode<T>(op->name(), ng_input, ng_reduction_indices, keep_dims);
+  SaveNgOp(ng_op_map, op->name(), ng_node);
+  return Status::OK();
 }
 
 static Status TranslateOneHotOp(
@@ -3908,8 +3877,8 @@ const static std::map<
         {"Add", TranslateBinaryOp<ngraph::opset3::Add>},
         {"AddN", TranslateAddNOp},
         {"AddV2", TranslateBinaryOp<ngraph::opset3::Add>},
-        {"Any", TranslateDirectReduceOp<ng::opset3::ReduceLogicalOr>},
-        {"All", TranslateDirectReduceOp<ng::opset3::ReduceLogicalAnd>},
+        {"Any", TranslateReduceOp<ng::opset3::ReduceLogicalOr>},
+        {"All", TranslateReduceOp<ng::opset3::ReduceLogicalAnd>},
         {"ArgMax", TranslateArgMinMaxOp<ng::op::ArgMax>},
         {"ArgMin", TranslateArgMinMaxOp<ng::op::ArgMin>},
         {"Asin", TranslateUnaryOp<ngraph::opset3::Asin>},
@@ -3961,13 +3930,13 @@ const static std::map<
         {"LogicalNot", TranslateUnaryOp<ngraph::opset3::LogicalNot>},
         {"LogicalOr", TranslateBinaryOp<ngraph::opset3::LogicalOr>},
         {"MatMul", TranslateMatMulOp},
-        {"Max", TranslateDirectReduceOp<ng::opset3::ReduceMax>},
+        {"Max", TranslateReduceOp<ng::opset3::ReduceMax>},
         {"Maximum", TranslateBinaryOp<ngraph::opset3::Maximum>},
         {"MaxPool", TranslateMaxPoolOp},
         {"MaxPool3D", TranslateMaxPool3DOp},
         {"NonMaxSuppressionV4", TranslateNonMaxSuppressionV4Op},
-        {"Mean", TranslateDirectReduceOp<ng::opset3::ReduceMean>},
-        {"Min", TranslateDirectReduceOp<ng::opset3::ReduceMin>},
+        {"Mean", TranslateReduceOp<ng::opset3::ReduceMean>},
+        {"Min", TranslateReduceOp<ng::opset3::ReduceMin>},
         {"Minimum", TranslateBinaryOp<ngraph::opset3::Minimum>},
         {"MirrorPad", TranslatePadOp},
         {"Mul", TranslateBinaryOp<ngraph::opset3::Multiply>},
@@ -3985,7 +3954,7 @@ const static std::map<
         {"Pow", TranslateBinaryOp<ngraph::opset3::Power>},
         // PreventGradient is just Identity in data-flow terms, so reuse that.
         {"PreventGradient", TranslateIdentityOp},
-        {"Prod", TranslateDirectReduceOp<ng::opset3::ReduceProd>},
+        {"Prod", TranslateReduceOp<ng::opset3::ReduceProd>},
         {"QuantizeAndDequantizeV2", TranslateQuantizeAndDequantizeV2Op},
         {"QuantizedAvgPool", TranslateQuantizedAvgPoolOp},
         {"QuantizedConcat", TranslateQuantizedConcatOp},
@@ -4029,7 +3998,7 @@ const static std::map<
         {"Squeeze", TranslateSqueezeOp},
         {"StridedSlice", TranslateStridedSliceOp},
         {"Sub", TranslateBinaryOp<ngraph::opset3::Subtract>},
-        {"Sum", TranslateDirectReduceOp<ng::opset3::ReduceSum>},
+        {"Sum", TranslateReduceOp<ng::opset3::ReduceSum>},
         {"Tan", TranslateUnaryOp<ngraph::opset3::Tan>},
         {"Tanh", TranslateUnaryOp<ngraph::opset3::Tanh>},
         {"Tile", TranslateTileOp},

ngraph_bridge/ngraph_mark_for_clustering.cc

@@ -196,25 +196,18 @@ const std::map<std::string, SetAttributesFunction>& GetAttributeSetters() {
 
   if (!initialized) {
     // Set Additional Attributes (if any)
-    set_attributes_map["Any"] = SetStaticInputs({1});
-    set_attributes_map["All"] = SetStaticInputs({1});
     set_attributes_map["ArgMax"] = SetStaticInputs({1});
     set_attributes_map["ArgMin"] = SetStaticInputs({1});
     set_attributes_map["ConcatV2"] = SetStaticInputs({-1});
     set_attributes_map["Conv2DBackpropInput"] = SetStaticInputs({0});
     set_attributes_map["ExpandDims"] = SetStaticInputs({1});
     set_attributes_map["Fill"] = SetStaticInputs({0});
     set_attributes_map["GatherV2"] = SetStaticInputs({2});
-    set_attributes_map["Max"] = SetStaticInputs({1});
-    set_attributes_map["Mean"] = SetStaticInputs({1});
-    set_attributes_map["Min"] = SetStaticInputs({1});
     set_attributes_map["MirrorPad"] = SetStaticInputs({1});
     set_attributes_map["NonMaxSuppressionV4"] = SetStaticInputs({2, 3, 4});
     set_attributes_map["OneHot"] = SetStaticInputs({1});
     set_attributes_map["Pad"] = SetStaticInputs({1});
     set_attributes_map["PadV2"] = SetStaticInputs({1, 2});
-    set_attributes_map["Prod"] = SetStaticInputs({1});
-
     set_attributes_map["QuantizeAndDequantizeV2"] = SetStaticInputs({1, 2});
     set_attributes_map["QuantizedConcat"] = [](Node* n) {
       SetStaticInputs(n, {0});  // the axis
@@ -242,7 +235,6 @@ const std::map<std::string, SetAttributesFunction>& GetAttributeSetters() {
     set_attributes_map["Split"] = SetStaticInputs({0});
     set_attributes_map["SplitV"] = SetStaticInputs({1, 2});
     set_attributes_map["StridedSlice"] = SetStaticInputs({1, 2, 3});
-    set_attributes_map["Sum"] = SetStaticInputs({1});
     set_attributes_map["TopKV2"] = SetStaticInputs({1});
     set_attributes_map["Tile"] = SetStaticInputs({1});
     set_attributes_map["Transpose"] = SetStaticInputs({1});