Check inputs num for node matching (#1885)

Fix inputs num mismatch for node matching. --------- Signed-off-by: xuzhenqi <[email protected]> Co-authored-by: xuzhenqi <[email protected]> Co-authored-by: Justin Chu <[email protected]>
microsoft · Oct 9, 2024 · a7c797d · a7c797d
1 parent 37b11fc
commit a7c797d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/onnxscript/rewriter/pattern.py b/onnxscript/rewriter/pattern.py
@@ -959,6 +959,12 @@ def _match_node(self, pattern_node: NodePattern, node: ir.Node) -> bool:
 
         self._matched[pattern_node] = node
 
+        # TODO: Revisit this to handle optional trailing inputs better.
+        if len(node.inputs) != len(pattern_node.inputs):
+            return self.fail(
+                "Input nums mismatch. {len(node.inputs)} vs {len(pattern_node.inputs)}"
+            )
+
         for arg_value, arg_pattern in zip(node.inputs, pattern_node.inputs):
             # arg_pattern could be a Var, if it's the original arg.
             if arg_pattern is None: