Cleanup shape inference implementation

onnx/shape_inference/implementation.cc

@@ -55,14 +55,6 @@
   }
 #endif
   return ONNX_NAMESPACE::to_string(type.elem_type());
-
-  void UnknownOpError(const NodeProto& nodeProto) {
-    if (checker::check_is_experimental_op(nodeProto)) {
-      fail_type_inference("Experimental operator '", nodeProto.op_type(), "' no longer upported.");
-    }
-    fail_type_inference(
-        "Unknown operator/function: '", nodeProto.op_type(), "' in domain: '", nodeProto.domain(), "'.");
-  }
 }
 
 std::string GetElemTypeString(const TypeProto_SparseTensor& type) {
@@ -78,6 +70,13 @@
 inline bool IsOnnxDomainOp(const NodeProto& node, const std::string& op_type) {
   return (IsOnnxDomain(node.domain()) && (node.op_type() == op_type));
 }
+
+void UnknownOpError(const NodeProto& nodeProto) {
+  if (checker::check_is_experimental_op(nodeProto)) {
+    fail_type_inference("Experimental operator '", nodeProto.op_type(), "' no longer upported.");
+  }
+  fail_type_inference("Unknown operator/function: '", nodeProto.op_type(), "' in domain: '", nodeProto.domain(), "'.");
+}
 } // namespace
 
 template <class T>
@@ -266,9 +265,9 @@
 // for inference. For GraphProto, inferred types are stored in the GraphProto
 // but FunctionProto does not have a place to store inferred types. So, we
 // use a temporary vector (for the duration of inference) to store these.
 class InferredTypes {
  public:
   InferredTypes(GraphProto* graph = nullptr) : graph_ptr(graph) {}
 
   TypeProto* Add(const std::string& var_name, const TypeProto& type) {
     if (graph_ptr != nullptr) {
@@ -298,7 +297,7 @@
 void BindValuesOnCall(
     const DataValueMap& caller_map,
     const NodeProto& caller,
     DataValueMap& callee_map,
     const FunctionProto& callee) {
   auto num_inputs = (std::min)(caller.input_size(), callee.input_size());
   for (int i = 0; i < num_inputs; ++i) {
@@ -317,7 +316,7 @@
 void BindValuesOnReturn(
     const DataValueMap& callee_map,
     const FunctionProto& callee,
     DataValueMap& caller_map,
     const NodeProto& caller) {
   auto num_outputs = (std::min)(caller.output_size(), callee.output_size());
   for (int i = 0; i < num_outputs; ++i) {
@@ -360,7 +359,7 @@
     }
   }
 
   void UpdateType(ValueInfoProto& valueInfo) {
     if (valueInfo.has_type()) {
       value_types_by_name[valueInfo.name()] = valueInfo.mutable_type();
     } else {
@@ -421,7 +420,7 @@
     }
   }
 
   void ProcessCall(const NodeProto& caller, const FunctionProto& callee, InferenceContext& ctx) {
     DataValueMap callee_value_map;
     if (generated_shape_data_by_name != nullptr) {
       BindValuesOnCall(*generated_shape_data_by_name, caller, callee_value_map, callee);
@@ -433,7 +432,7 @@
     }
   }
 
   void Process(NodeProto& n) {
     // Resolve domain for node
     auto dit = opset_imports.find(n.domain());
     if (dit == opset_imports.end()) {
@@ -470,9 +469,13 @@
           ProcessCall(n, *(schema->GetFunction()), ctx);
         } else {
           // Continue with inference for remaining nodes
           // TODO: fix this
           return;
         }
+        // check type-constraints specified via type variables
+        if (options.check_type) {
+          schema->CheckInputOutputType(ctx);
+        }
       } else if (model_local_functions_map.size() > 0) {
         auto iter = model_local_functions_map.find(GetModelLocalFunctionsMapIdentifier(n.domain(), n.op_type()));
         if (iter != model_local_functions_map.end()) {
@@ -483,28 +486,15 @@
       } else {
         UnknownOpError(n);
       }
-    }
-    ONNX_CATCH(const ONNX_NAMESPACE::InferenceError& ex) {
-      ONNX_HANDLE_EXCEPTION([&]() { inference_errors.push_back(GetErrorWithNodeInfo(n, ex)); });
-      // Continue with inference for remaining nodes
-      return;
-    }
-
-    ONNX_TRY {
-      // check the type-equality for input and output
-      if (options.check_type && schema) {
-        schema->CheckInputOutputType(ctx);
-      }
-
       for (int i = 0; i < n.output_size(); ++i) {
         // skip type and shape propagation for missing optional outputs.
         if (!n.output(i).empty())
           UpdateType(n.output(i), ctx.getOutputType(i));
       }
-
+      // Constant values are tracked to improve inference/checking for subsequent nodes.
       ProcessConstant(n);
-
-      // If data propagation is enabled, propagate shape data if it exists.
+      // If data-propagation is enabled, partial-evaluation (aka data-propagation) is performed
+      // to improve inference/checking for subsequent nodes.
       if (options.enable_data_propagation && schema && schema->has_data_propagation_function()) {
         if (generated_shape_data_by_name == nullptr) {
           fail_shape_inference(
@@ -515,7 +505,11 @@
         schema->GetDataPropagationFunction()(data_propagation_ctx);
       }
     }
+    ONNX_CATCH(const ONNX_NAMESPACE::InferenceError& ex) {
+      ONNX_HANDLE_EXCEPTION([&]() { inference_errors.push_back(GetErrorWithNodeInfo(n, ex)); });
+    }
     ONNX_CATCH(const std::runtime_error& err) {
+      // TODO: Fix this. Unclear if this should be remapped to a shape inference error.
       ONNX_HANDLE_EXCEPTION([&]() { fail_shape_inference(GetErrorWithNodeInfo(n, err)); });
     }
   }
@@ -544,7 +538,7 @@
     }
   }
 
   void Process(GraphProto& graph) {
     if (symbol_table) {
       TraverseGraphsToAddExistingSymbols(graph, *symbol_table);
     }
@@ -584,7 +578,7 @@
     }
   }
 
   void Process(const NodeProto& n, internal::AttributeBinder& attribute_binder) {
     NodeProto copy_n(n);
     attribute_binder.VisitNode(&copy_n);
     Process(copy_n);
@@ -695,7 +689,7 @@
   void FinalizeShapeInference() {
     auto& errors = getErrors();
     // Throw shape inference error if any. Error mode right now only supports 0 and 1.
     // When set to 0, any node level shape inference errors are not thrown. This is to support backward compatiblity
     // with 1.7 and earlier releases. When set to 1 it will throw all exceptions.
     // TODO: Add a more granular way for exception handling.
     if (!errors.empty() && options.error_mode > 0) {