QuadricCustomOp: Handle multiple outputs when shape inferencing (#17)

quadric-io · Oct 3, 2023 · 126cce8 · 126cce8
1 parent d0776ea
commit 126cce8
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
@@ -2930,7 +2930,6 @@ This op functions in much the same was as Dropout-11 and Dropout-13 do, execpt t
       .Attr("element_wise", "True (1) if only element-wise ops, False (0) otherwise", AttributeProto::INT, true)
       .TypeConstraint("T", OpSchema::all_tensor_types_with_bfloat(),
                       "Allow inputs and outputs to be any kind of tensor.");
-  // FIXME: Add a type/shape inference function
 
 #ifdef ENABLE_TRAINING_OPS
   // Should remove the shrunken_gather include from ENABLE_TRAINING_OPS once 1). compute optimizer is enabled for inference or

diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py
@@ -226,7 +226,6 @@ def __init__(self, int_max, auto_merge, guess_output_rank, verbose, prefix=""):
             "QLinearAveragePool": self._infer_qlinear_unary_op,
             # Quadric custom operators
             "QuadricCustomOp": self._infer_custom_op,
-            "QuadricCustomOpElementWise": self._infer_custom_op
         }
         self.aten_op_dispatcher_ = {
             "embedding": self._infer_Gather,
@@ -455,6 +454,7 @@ def _onnx_infer_single_node(self, node):
             "If",
             "Loop",
             "Scan",
+            "QuadricCustomOp",
             "SplitToSequence",
             "ZipMap",  # contrib ops
             "Attention",
@@ -974,17 +974,21 @@ def _infer_qgemm(self, node):
     def _infer_custom_op(self, node):
         # For the CCL custom operators the shape and dtype of the output are present in
         # the attributes and can be used to directly create the value info
-        attr_map = {n.name:n for n in list(node.attribute)}
-        assert "shape" in attr_map and "elem_type" in attr_map,\
-            "Custom op output type not found"
-        vi = self.known_vi_[node.output[0]]
-        vi.CopyFrom(
-            helper.make_tensor_value_info(
-                    node.output[0],
-                    attr_map["elem_type"].i,
-                    attr_map["shape"].ints,
+        attr_map = {n.name: n for n in list(node.attribute)}
+        assert "shape" in attr_map and "elem_type" in attr_map, "Custom op output type not found"
+        if len(node.output) > 1:
+            for i, out in enumerate(node.output):
+                vi = self.known_vi_[out]
+                vi.CopyFrom(
+                    helper.make_tensor_value_info(
+                        out,
+                        attr_map["elem_type"].ints[i],
+                        attr_map["shape"].tensors[i].int32_data,
+                    )
                 )
-            )
+        else:
+            vi = self.known_vi_[node.output[0]]
+            vi.CopyFrom(helper.make_tensor_value_info(node.output[0], attr_map["elem_type"].i, attr_map["shape"].ints))
 
     def _infer_ConcatFromSequence(self, node):
         seq_shape = self._get_shape(node, 0)
@@ -2582,6 +2586,10 @@ def get_prereq(node):
                     get_attribute(node, "then_branch"),
                     get_attribute(node, "else_branch"),
                 ]
+            elif node.op_type == "QuadricCustomOp":
+                # Should have a subgraph, but allow for cases where it's not there
+                subgraph = get_attribute(node, "sub_graph")
+                subgraphs = [subgraph] if subgraph else []
             elif node.op_type in ["Loop", "Scan"]:
                 subgraphs = [get_attribute(node, "body")]
             for g in subgraphs: