fix ut

Author: xadupre

_doc/examples/plot_export_tiny_llm.py

@@ -33,6 +33,7 @@
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.torch_helper import steal_forward
 from onnx_diagnostic.torch_models.llms import get_tiny_llm
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 
 
 MODEL_NAME = "arnir0/Tiny-LLM"
@@ -131,7 +132,11 @@ def _forward_(*args, _f=None, **kwargs):
 
 try:
     ep = torch.export.export(
-        untrained_model, (), kwargs=cloned_inputs, dynamic_shapes=dynamic_shapes, strict=False
+        untrained_model,
+        (),
+        kwargs=cloned_inputs,
+        dynamic_shapes=use_dyn_not_str(dynamic_shapes),
+        strict=False,
     )
     print("It worked:")
     print(ep)
@@ -166,7 +171,11 @@ def _forward_(*args, _f=None, **kwargs):
 
 try:
     ep = torch.export.export(
-        model, (), kwargs=cloned_inputs, dynamic_shapes=dynamic_shapes, strict=False
+        model,
+        (),
+        kwargs=cloned_inputs,
+        dynamic_shapes=use_dyn_not_str(dynamic_shapes),
+        strict=False,
     )
     print("It worked:")
     print(ep)

_unittests/ut_torch_models/test_tiny_llms_bypassed.py

@@ -2,19 +2,21 @@
 import unittest
 import torch
 from transformers.cache_utils import DynamicCache
-from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings
+from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings, hide_stdout
 from onnx_diagnostic.torch_models.llms import get_tiny_llm
 from onnx_diagnostic.torch_models.llms import get_phi2
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.torch_helper import torch_deepcopy
 from onnx_diagnostic.torch_export_patches import torch_export_patches
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 from onnx_diagnostic.torch_export_patches.patches.patch_transformers import (
     patched_DynamicCache,
 )
 
 
 class TestTinyLlmBypassed(ExtTestCase):
     @ignore_warnings(UserWarning)
+    @hide_stdout()
     def test_export_tiny_llm_2_bypassed(self):
         data = get_tiny_llm()
         model, inputs = data["model"], data["inputs"]
@@ -50,7 +52,11 @@ def debug():
                 debug()
 
             ep = torch.export.export(
-                model, (), kwargs=inputs, dynamic_shapes=data["dynamic_shapes"], strict=False
+                model,
+                (),
+                kwargs=inputs,
+                dynamic_shapes=use_dyn_not_str(data["dynamic_shapes"]),
+                strict=False,
             )
             got = ep.module()(**inputs)
             self.assertEqualArrayAny(expected, got)