fix make_static_cache

Author: xadupre

.github/workflows/check-release.yml

@@ -16,7 +16,7 @@ jobs:
       matrix:
         os: [ubuntu-latest, macOS-latest, windows-latest]
         python: ['3.11', '3.12']
-        transformers: ['4.48.3', '4.52.4', 'main']
+        transformers: ['4.48.3', '4.52.4', '4.55.2', 'main']
         torch: ['2.7', '2.8', 'main']
 
     steps:

.github/workflows/ci.yml

@@ -16,7 +16,7 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         python: ['3.10', '3.11', '3.12', '3.13']
-        transformers: ['4.48.3', '4.51.3', '4.52.4', '4.53.3', '4.55.0', 'main']
+        transformers: ['4.48.3', '4.51.3', '4.52.4', '4.53.3', '4.55.2', 'main']
         torch: ['2.8', 'main']
         exclude:
           - python: '3.10'
@@ -28,15 +28,15 @@ jobs:
           - python: '3.10'
             transformers: '4.53.3'
           - python: '3.10'
-            transformers: '4.55.0'
+            transformers: '4.55.2'
           - python: '3.11'
             torch: 'main'
           - python: '3.11'
             transformers: '4.53.3'
           - python: '3.11'
             transformers: 'main'
           - python: '3.11'
-            transformers: '4.55.0'
+            transformers: '4.55.2'
           - python: '3.13'
             torch: '2.8'
           - python: '3.13'

onnx_diagnostic/helpers/cache_helper.py

@@ -280,6 +280,33 @@ def __init__(self):
         max_cache_len=max_cache_len,
     )
     ca = CacheKeyValue(cache)
+    if hasattr(cache, "layers") and len(ca.key_cache) == 0:
+        # transformers>= 4.55.2, layers are empty
+        for i, (key, value) in enumerate(key_value_pairs):
+            cache.update(key, value, i)
+        return cache
+
+    torch._check(
+        len(key_value_pairs) == len(cache.layers),
+        lambda: (
+            f"Length mismatch len(key_value_pairs)={len(key_value_pairs)}, "
+            f"len(cache.layers)={len(cache.layers)}"
+        ),
+    )
+    torch._check(
+        len(key_value_pairs) == len(ca.key_cache),
+        lambda: (
+            f"Length mismatch len(key_value_pairs)={len(key_value_pairs)}, "
+            f"len(ca.key_cache)={len(ca.key_cache)}"
+        ),
+    )
+    torch._check(
+        len(key_value_pairs) == len(ca.value_cache),
+        lambda: (
+            f"Length mismatch len(key_value_pairs)={len(key_value_pairs)}, "
+            f"len(ca.value_cache)={len(ca.value_cache)}"
+        ),
+    )
     for i in range(len(key_value_pairs)):
         assert (
             key_value_pairs[i][0].shape == key_value_pairs[i][1].shape

requirements.txt

@@ -2,5 +2,5 @@ numpy
 onnx>=1.16.0
 onnxruntime>=1.21
 optree
-torch>=2.7
+torch>=2.8
 torch_geometric