fix static

Author: xadupre

_unittests/ut_torch_models/test_tiny_llms.py

@@ -1,7 +1,7 @@
 import copy
 import unittest
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings
+from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings, requires_transformers
 from onnx_diagnostic.torch_models.llms import get_tiny_llm
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_export_patches import torch_export_patches
@@ -33,13 +33,15 @@ def test_tiny_llm_export_dynamic(self):
             got = ep.module()(**inputs)
             self.assertEqualArrayAny(expected, got)
 
+    @requires_transformers("4.52")
     def test_tiny_llm_run_static(self):
         data = get_tiny_llm(use_static_cache=True)
         model, inputs = data["model"], data["inputs"]
         self.assertIn("StaticCache", string_type(inputs))
         model(**inputs)
 
     @ignore_warnings(UserWarning)
+    @requires_transformers("4.52")
     def test_tiny_llm_export_static(self):
         data = get_tiny_llm(use_static_cache=True)
         model, inputs = data["model"], data["inputs"]

onnx_diagnostic/torch_models/untrained/llm_tiny_llm.py

@@ -50,7 +50,8 @@ def get_tiny_llm(
 
     config.update(**kwargs)
     conf = transformers.LlamaConfig(**config)
-    conf.cache_implementation = "static"
+    if use_static_cache:
+        conf.cache_implementation = "static"
     model = transformers.LlamaForCausalLM(conf)
     model.eval()