Fix test

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

src/onnx_ir/passes/common/clear_metadata_and_docstring_test.py

@@ -35,20 +35,36 @@ def test_pass_with_clear_metadata_and_docstring(self):
             metadata_props={"mul_key": "mul_value"},
             doc_string="This is a Mul node",
         )
-        func_inputs = [
-            ir.Value(
-                name="input_a", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3))
-            ),
-            ir.Value(
-                name="input_b", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3))
-            ),
-        ]
         function = ir.Function(
             graph=ir.Graph(
                 name="my_function",
-                inputs=func_inputs,
-                outputs=mul_node.outputs,
-                nodes=[add_node, mul_node],
+                inputs=[
+                    input_a := ir.Value(
+                        name="input_a",
+                        type=ir.TensorType(ir.DataType.FLOAT),
+                        shape=ir.Shape((2, 3)),
+                    ),
+                    input_b := ir.Value(
+                        name="input_b",
+                        type=ir.TensorType(ir.DataType.FLOAT),
+                        shape=ir.Shape((2, 3)),
+                    ),
+                ],
+                nodes=[
+                    add_node_func := ir.node(
+                        "Add",
+                        inputs=[input_a, input_b],
+                        metadata_props={"add_key": "add_value"},
+                        doc_string="This is an Add node",
+                    ),
+                    mul_node_func := ir.node(
+                        "Mul",
+                        inputs=[add_node_func.o(), input_b],
+                        metadata_props={"mul_key": "mul_value"},
+                        doc_string="This is a Mul node",
+                    ),
+                ],
+                outputs=mul_node_func.outputs,
                 opset_imports={"": 20},
                 doc_string="This is a function docstring",
                 metadata_props={"function_key": "function_value"},
@@ -59,8 +75,8 @@ def test_pass_with_clear_metadata_and_docstring(self):
         )
         func_node = ir.node(
             "my_function",
-            inputs=[add_node.o(), inputs[1]],
-            domain = "my_domain",
+            inputs=[inputs[0], mul_node.o()],
+            domain="my_domain",
             metadata_props={"mul_key": "mul_value"},
             doc_string="This is a Mul node",
         )
@@ -77,7 +93,7 @@ def test_pass_with_clear_metadata_and_docstring(self):
         )
         sub_node = ir.node(
             "Sub",
-            inputs=[function.o(), const_node.o()],
+            inputs=[func_node.o(), const_node.o()],
             num_outputs=1,
             metadata_props={"sub_key": "sub_value"},
             doc_string="This is a Sub node",