fix

justinchuby · justinchuby · commit b99fe8896029 · 2025-09-03T12:13:08.000-07:00
Signed-off-by: Justin Chu &lt;justinchuby@users.noreply.github.com&gt;
diff --git a/onnxscript/optimizer/_constant_folding_test.py b/onnxscript/optimizer/_constant_folding_test.py
@@ -599,7 +599,7 @@ def test_node_is_folded_if_specified_as_should_fold(self):
 
         # But ConstantOfShape is folded when specified in should_fold
         optimized = self._fold(
-            model, should_fold=lambda node: node.op_type == "ConstantOfShape"
+            model, should_fold=lambda node: node.op_type == "ConstantOfShape" or None
         )
         ops = [node.op_type for node in optimized.graph]
         self.assertEqual(ops, ["Constant"])

Original file line number	Diff line number	Diff line change
`@@ -599,7 +599,7 @@ def test_node_is_folded_if_specified_as_should_fold(self):`
`599`	`599`
`600`	`600`	`# But ConstantOfShape is folded when specified in should_fold`
`601`	`601`	`optimized = self._fold(`
`602`		`- model, should_fold=lambda node: node.op_type == "ConstantOfShape"`
	`602`	`+ model, should_fold=lambda node: node.op_type == "ConstantOfShape" or None`
`603`	`603`	`)`
`604`	`604`	`ops = [node.op_type for node in optimized.graph]`
`605`	`605`	`self.assertEqual(ops, ["Constant"])`