Loop type shape inferencing (onnx#1591)

* Add Loop type/shape inferencing.
Make Loop spec more consistent.

* Remove temporary debug code.

* Make work with current Loop spec.

* Allow loop carried dependencies to change shape across iterations.

* Make iter_num_in type correct.

* Check iteration num type is tensor(int64) to match the max iterations input.

Update mergeShapesAndTypes to fail if the types don't match.

Author: skottmckay

Diff for: onnx/defs/controlflow/defs.cc

@@ -272,6 +272,117 @@ void IfInferenceFunction(InferenceContext& ctx) {
   }
 }
 
+void LoopInferenceFunction(InferenceContext& ctx) {
+  auto num_inputs = ctx.getNumInputs();
+  auto num_loop_state_vars = num_inputs - 2; // skip 'M' and 'cond'
+
+  std::vector<const TypeProto*> subgraph_input_types;
+
+  std::vector<TypeProto> temporary_type_protos;
+  temporary_type_protos.reserve(num_inputs - 2);
+
+  // create TypeProto to validate iteration number type is the same as the
+  // optional 'M' input for max iterations.
+  TypeProto iter_num_type;
+  iter_num_type.mutable_tensor_type()->set_elem_type(
+      TensorProto_DataType_INT64);
+  subgraph_input_types.push_back(&iter_num_type);
+
+  // 'cond'
+  subgraph_input_types.push_back(ctx.getInputType(1));
+
+  // loop state value types get propagated to outputs, but shape may change
+  // across iterations so don't propagate it to the outputs and don't pass it
+  // into the subgraph inferencing
+  for (size_t i = 2; i < num_inputs; ++i) {
+    propagateElemTypeFromInputToOutput(ctx, i, i - 2);
+
+    // copy so we can remove the shape before passing to the subgraph
+    // inferencing
+    temporary_type_protos.push_back(*ctx.getInputType(i));
+    auto& input_type = temporary_type_protos.back();
+    input_type.mutable_tensor_type()->clear_shape();
+
+    subgraph_input_types.push_back(&input_type);
+  }
+
+  // Run inferencing on the subgraph
+  std::vector<const TypeProto*> subgraph_output_types;
+
+  GraphInferencer* graphInferencer = ctx.getGraphAttributeInferencer("body");
+  if (graphInferencer) {
+    std::vector<const TensorProto*> input_data;
+    input_data.push_back(nullptr); // iteration number
+    for (size_t i = 1; i < num_inputs; ++i) {
+      input_data.push_back(ctx.getInputData(i));
+    }
+
+    subgraph_output_types =
+        graphInferencer->doInferencing(subgraph_input_types, input_data);
+  }
+
+  // if empty(), assume inferencing was skipped
+  if (!subgraph_output_types.empty()) {
+    auto num_outputs = ctx.getNumOutputs();
+
+    // subgraph outputs the condition value first but that is only used
+    // internally and not returned by Loop.
+    if (subgraph_output_types.size() != num_outputs + 1) {
+      fail_type_inference(
+          "Graph attribute inferencing returned type information for ",
+          subgraph_output_types.size(),
+          " outputs. Expected ",
+          num_outputs + 1);
+    }
+
+    // check loop state values match. we should already have type/shape info
+    for (size_t i = 0; i < num_outputs; ++i) {
+      auto* subgraph_output_type = subgraph_output_types[i + 1]; // skip 'cond'
+      auto* loop_output_type = ctx.getOutputType(i);
+
+      const bool is_loop_state_var = i < num_loop_state_vars;
+
+      if (!subgraph_output_type->has_tensor_type()) {
+        fail_type_inference(
+            "Loop 'body' subgraph outputs should all be tensors but output ",
+            i,
+            " was ",
+            subgraph_output_type->value_case());
+      }
+
+      // if there's an existing type check it matches. otherwise propagate
+      propagateElemTypeWithValidation(subgraph_output_type, loop_output_type);
+
+      if (is_loop_state_var) {
+        // shape may change across iterations so ignore.
+      } else {
+        // per iteration output. first dimension will be number of iterations
+        // but we don't know that value yet
+        TypeProto inferred_type(*subgraph_output_type);
+        auto* mutable_inferred_tensor_type =
+            inferred_type.mutable_tensor_type();
+        auto* mutable_inferred_shape =
+            mutable_inferred_tensor_type->mutable_shape();
+
+        mutable_inferred_shape->clear_dim();
+
+        // add empty dimension for number of iterations
+        mutable_inferred_shape->add_dim();
+
+        // add dimensions from subgraph output shape
+        for (const auto& dim :
+             subgraph_output_type->tensor_type().shape().dim()) {
+          (*mutable_inferred_shape->add_dim()) = dim;
+        }
+
+        mergeInShapeInfo(
+            *mutable_inferred_tensor_type,
+            *loop_output_type->mutable_tensor_type());
+      }
+    }
+  }
+}
+
 ONNX_OPERATOR_SET_SCHEMA(
     If,
     1,
@@ -459,7 +570,8 @@ ONNX_OPERATOR_SET_SCHEMA(
             AttributeProto::GRAPH)
         .TypeConstraint("V", OpSchema::all_tensor_types(), "All Tensor types")
         .TypeConstraint("I", {"int64"}, "Only int64")
-        .TypeConstraint("B", {"bool"}, "Only bool"));
+        .TypeConstraint("B", {"bool"}, "Only bool")
+        .TypeAndShapeInferenceFunction(LoopInferenceFunction));
 
 static const char* scan_ver1_doc = R"DOC(
 Scan can be used to iterate over one or more scan_input tensors,

Diff for: onnx/defs/shape_inference.h

@@ -151,6 +151,50 @@ multiplyDims(const TensorShapeProto& shape, int from, int upto_exclusive) {
   return dim;
 }
 
+// propagate the element type from an input type to an output type.
+// if an existing output element type exists, validate it matches.
+inline void propagateElemTypeWithValidation(
+    const TypeProto* input_type,
+    TypeProto* output_type) {
+  if (nullptr == input_type) {
+    fail_type_inference("Input type was null");
+  }
+
+  if (input_type->value_case() != TypeProto::kTensorType) {
+    fail_type_inference(
+        "Input was expected to have tensor type. Got ",
+        input_type->value_case());
+  }
+
+  if (input_type->tensor_type().elem_type() == TensorProto::UNDEFINED) {
+    fail_type_inference("Element type of input was unknown");
+  }
+
+  if (output_type->value_case() == TypeProto::VALUE_NOT_SET) {
+    output_type->mutable_tensor_type()->set_elem_type(
+        input_type->tensor_type().elem_type());
+  } else if (output_type->value_case() == TypeProto::kTensorType) {
+    if (output_type->tensor_type().has_elem_type()) {
+      if (input_type->tensor_type().elem_type() !=
+          output_type->tensor_type().elem_type()) {
+        fail_type_inference(
+            "Input element type of ",
+            input_type->tensor_type().elem_type(),
+            " does not match existing output type of ",
+            output_type->tensor_type().elem_type());
+      }
+    } else {
+      output_type->mutable_tensor_type()->set_elem_type(
+          input_type->tensor_type().elem_type());
+    }
+  } else {
+    // This is not expected to happen
+    fail_type_inference(
+        "Output was expected to have tensor type. Got ",
+        output_type->value_case());
+  }
+}
+
 // Note: for all methods below for propagating type or shape, callers are
 // responsible to handle optional inputs/outputs and ensure that the specified
 // index value is less than NumInputs/NumOutputs.

Diff for: onnx/shape_inference/implementation.cc

@@ -44,9 +44,16 @@ void checkShapesAndTypes(
 void mergeShapesAndTypes(
     const TypeProto_Tensor& inferredType,
     TypeProto_Tensor* existingType) {
-  if (inferredType.elem_type() != TensorProto::UNDEFINED &&
-      existingType->elem_type() == TensorProto::UNDEFINED) {
-    existingType->set_elem_type(inferredType.elem_type());
+  if (inferredType.elem_type() != TensorProto::UNDEFINED) {
+    if (existingType->elem_type() == TensorProto::UNDEFINED) {
+      existingType->set_elem_type(inferredType.elem_type());
+    } else if (existingType->elem_type() != inferredType.elem_type()) {
+      fail_type_inference(
+          "type mismatch. existing=",
+          existingType->elem_type(),
+          " inferred=",
+          inferredType.elem_type());
+    }
   }
 
   if (!inferredType.has_shape()) {
@@ -324,6 +331,10 @@ std::vector<const TypeProto*> GraphInferencerImpl::doInferencing(
 
   for (int i = 0, end = numInputs; i < end; ++i) {
     const TypeProto* inferredInput = inputTypes[i];
+
+    if (!inferredInput)
+      continue;
+
     TypeProto* graphInput = g_->mutable_input(i)->mutable_type();
 
     if (!graphInput->has_tensor_type()) {

Diff for: onnx/test/shape_inference_test.py

@@ -956,9 +956,6 @@ def test_scan(self):    # type: () -> None
         # can't use self._make_graph for the subgraph as it add more inputs for the Reshape operations it inserts.
         # this breaks the subgraph inferencing as it expects the number of inputs passed from Scan to match
         # the GraphProto, but Scan knows nothing about the additional inputs.
-        value_infos = [make_tensor_value_info('loop_state_in', TensorProto.FLOAT, (loop_state_size,)),
-                       make_tensor_value_info('input', TensorProto.FLOAT, (input_size,))]
-
         input_value_infos = [make_tensor_value_info('loop_state_in', TensorProto.UNDEFINED, None),
                              make_tensor_value_info('input', TensorProto.UNDEFINED, None)]
         output_value_infos = [make_tensor_value_info('loop_state_out', TensorProto.UNDEFINED, None),
@@ -969,15 +966,14 @@ def test_scan(self):    # type: () -> None
              make_node('Identity', ['input'], ['output'])],
             "subgraph",
             input_value_infos,
-            output_value_infos,
-            value_info=value_infos
+            output_value_infos
         )
 
         graph = self._make_graph(
             [('loop_state_orig', TensorProto.FLOAT, (batch_size, loop_state_size)),
              ('scan_input', TensorProto.FLOAT, (batch_size, seq_len, input_size))],
             [make_node('Scan', ['', 'loop_state_orig', 'scan_input'], ['loop_state_final', 'scan_output'],
-                               num_scan_inputs=1, body=subgraph)],
+                       num_scan_inputs=1, body=subgraph)],
             []
         )
 
@@ -1053,6 +1049,41 @@ def test_onehot_with_axis(self):  # type: () -> None
             [])
         self._assert_inferred(graph, [make_tensor_value_info('Y', TensorProto.FLOAT, (2, None, 3, 5))])  # type: ignore
 
+    def test_loop(self):    # type: () -> None
+        # can't use self._make_graph for the subgraph as it add more inputs for the Reshape operations it inserts.
+        # this breaks the subgraph inferencing as it expects the number of inputs passed from Loop to match
+        # the GraphProto, but Loop knows nothing about the additional inputs.
+        input_value_infos = [make_tensor_value_info('iter_num_in', TensorProto.INT64, (1,)),
+                             make_tensor_value_info('cond_in', TensorProto.UNDEFINED, None),
+                             make_tensor_value_info('loop_state_in', TensorProto.UNDEFINED, ())]
+        output_value_infos = [make_tensor_value_info('cond_out', TensorProto.UNDEFINED, None),
+                              make_tensor_value_info('loop_state_out', TensorProto.UNDEFINED, None),
+                              make_tensor_value_info('output', TensorProto.FLOAT, (3,))]
+
+        subgraph = helper.make_graph(
+            [make_node('Identity', ['cond_in'], ['cond_out']),
+             make_node('Identity', ['loop_state_in'], ['loop_state_out']),
+             make_node('Identity', ['outer_scope_input'], ['output'])],
+            "subgraph",
+            input_value_infos,
+            output_value_infos
+        )
+
+        graph = self._make_graph(
+            [('max_trip_count', TensorProto.INT64, (1,)),
+             ('cond_orig', TensorProto.FLOAT, (1,)),
+             ('loop_state_orig', TensorProto.FLOAT, (2,)),
+             ('outer_scope_input', TensorProto.FLOAT, (3,))],
+            [make_node('Loop', ['max_trip_count', 'cond_orig', 'loop_state_orig'], ['loop_state_final', 'loop_output'],
+                       body=subgraph)],
+            []
+        )
+
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info('loop_state_final', TensorProto.FLOAT, None),  # shape may change between iterations
+             make_tensor_value_info('loop_output', TensorProto.FLOAT, (None, 3))])  # type: ignore
+
 
 if __name__ == '__main__':
     unittest.main()