mypy

Author: xadupre

_unittests/ut_torch_models/test_tiny_llms_onnx.py

@@ -28,14 +28,15 @@ def test_onnx_export_tiny_llm_official(self):
         self.assertEqual(
             {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
         )
-        ep = torch.onnx.export(
-            model,
-            (),
-            kwargs=inputs,
-            dynamic_shapes=data["dynamic_shapes"],
-            dynamo=True,
-            optimize=True,
-        )
+        with torch_export_patches(patch_transformers=True):
+            ep = torch.onnx.export(
+                model,
+                (),
+                kwargs=inputs,
+                dynamic_shapes=data["dynamic_shapes"],
+                dynamo=True,
+                optimize=True,
+            )
         # There are some discrepancies with torch==2.6
         if not has_torch("2.7"):
             raise unittest.SkipTest("discrepancies observed with torch<2.7")

onnx_diagnostic/helpers/graph_helper.py

@@ -290,9 +290,9 @@ def text_rendering(self, prefix="") -> str:
             *self.nodes,
             *[oh.make_node(i, [i], ["END"]) for i in self.output_names],
         ]
-        existing = set(self.start_names) - set(self.input_names)
-        existing |= {"BEGIN"}
-        existing = sorted(existing)
+        exist = set(self.start_names) - set(self.input_names)
+        exist |= {"BEGIN"}
+        existing = sorted(exist)
         order = self.computation_order(nodes, existing)
         positions = self.graph_positions(nodes, order, existing)
         text_pos = self.text_positions(nodes, positions)