fix split-18

daquexian · daquexian · commit 74fdf9cbbf9f · 2022-11-15T17:35:00.000+08:00
Signed-off-by: daquexian &lt;daquexian566@gmail.com&gt;
diff --git a/onnxoptimizer/passes/eliminate_nop_split.h b/onnxoptimizer/passes/eliminate_nop_split.h
@@ -23,9 +23,10 @@ struct EliminateNopSplit final : public PredicateBasedPass {
   }
 
   bool patternMatchPredicate(Node* node) override {
-    return CheckKind(node, "Split") && node->input()->has_sizes() &&
-           node->hasAttribute(kaxis) && node->hasAttribute(ksplit) &&
-           node->outputs().size() == 1;
+    // TODO: support Split-18 where `split` is an input instead of an attribute
+    return CheckKind(node, "Split") && node->inputs().size() == 1 &&
+           node->input()->has_sizes() && node->hasAttribute(kaxis) &&
+           node->hasAttribute(ksplit) && node->outputs().size() == 1;
   }
 
   bool runTransform(Node* node, Graph& graph,

Original file line number	Diff line number	Diff line change
`@@ -23,9 +23,10 @@ struct EliminateNopSplit final : public PredicateBasedPass {`
`23`	`23`	`}`
`24`	`24`
`25`	`25`	`bool patternMatchPredicate(Node* node) override {`
`26`		`- return CheckKind(node, "Split") && node->input()->has_sizes() &&`
`27`		`- node->hasAttribute(kaxis) && node->hasAttribute(ksplit) &&`
`28`		`- node->outputs().size() == 1;`
	`26`	+ // TODO: support Split-18 where `split` is an input instead of an attribute
	`27`	`+ return CheckKind(node, "Split") && node->inputs().size() == 1 &&`
	`28`	`+ node->input()->has_sizes() && node->hasAttribute(kaxis) &&`
	`29`	`+ node->hasAttribute(ksplit) && node->outputs().size() == 1;`
`29`	`30`	`}`
`30`	`31`
`31`	`32`	`bool runTransform(Node* node, Graph& graph,`