fix

Author: xadupre

_unittests/ut_investigate/test_input_observer.py

@@ -2,7 +2,12 @@
 import unittest
 import pandas
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, requires_torch
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    requires_torch,
+    hide_stdout,
+    ignore_warnings,
+)
 from onnx_diagnostic.investigate.input_observer import (
     InputObserver,
     _infer_dynamic_dimensions,
@@ -816,6 +821,8 @@ def forward(self, x=None, y=None):
         self.assertEqual(2, len(args))
         self.assertEqual(len([v for v in args.values() if v is not None]), 2)
 
+    @hide_stdout()
+    @ignore_warnings(FutureWarning)
     def test_io_int_kwargs(self):
         class Model(torch.nn.Module):
             def forward(self, x=None, y=None, option=1):

_unittests/ut_torch_export_patches/test_patch_transformers.py

@@ -12,6 +12,7 @@
     requires_torch,
     ignore_warnings,
     has_onnxscript,
+    requires_onnxscript,
 )
 from onnx_diagnostic.helpers.torch_helper import torch_deepcopy, fake_torchdynamo_exporting
 from onnx_diagnostic.export.shape_helper import make_fake_with_dynamic_dimensions
@@ -352,6 +353,7 @@ def forward(self, query, key, value):
         self.assertEqualArray(expected, got)
 
     @requires_transformers("4.55")
+    @requires_onnxscript("0.6.2")
     @unittest.skipIf(not patch_qwen2_5, "Qwen25 not part of this transformers")
     def test_qwen_apply_multimodal_rotary_pos_emb(self):
         apply_multimodal_rotary_pos_emb = (

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -162,20 +162,22 @@ def patched__get_range_constraints(
     combined_args = torch.export._trace._combine_args(mod, args, kwargs)
 
     # _combine_args does not preserve the order.
-    if isinstance(combined_args, dict):
-        input_names = [
-            s.arg.name
-            for s in export_graph_signature.input_specs
-            if s.kind == torch.export.graph_signature.InputKind.USER_INPUT
-        ]
-        new_args = {}
-        for k in input_names:
-            if k in combined_args:
-                new_args[k] = combined_args[k]
-        for k in combined_args:
-            if k not in new_args:
-                new_args[k] = combined_args[k]
-        combined_args = new_args
+    assert isinstance(
+        combined_args, dict
+    ), f"unexpected type {type(combined_args)} for 'combined_args'"
+    input_names = [
+        s.arg.name
+        for s in export_graph_signature.input_specs
+        if s.kind == torch.export.graph_signature.InputKind.USER_INPUT
+    ]
+    new_args = {}
+    for k in input_names:
+        if k in combined_args:
+            new_args[k] = combined_args[k]
+    for k in combined_args:
+        if k not in new_args:
+            new_args[k] = combined_args[k]
+    combined_args = new_args
 
     range_constraints = torch._export.non_strict_utils.make_constraints(
         fake_mode, gm, combined_args, dynamic_shapes, num_lifted