fix cache

Author: xadupre

onnx_diagnostic/ext_test_case.py

@@ -1014,7 +1014,7 @@ def assertEqualArrayAny(
                     msg_ = "\n".join(excs)
                     msg = f"{msg}\n{msg_}" if msg else msg_
                     raise AssertionError(f"Found {len(excs)} discrepancies\n{msg}")
-        elif expected.__class__.__name__ == "DynamicCache":
+        elif expected.__class__.__name__ in ("DynamicCache", "StaticCache"):
             atts = {"key_cache", "value_cache"}
             self.assertEqualArrayAny(
                 {k: expected.__dict__.get(k, None) for k in atts},

onnx_diagnostic/tasks/text_generation.py

@@ -174,7 +174,7 @@ def get_inputs(
             shapes = {
                 "input_ids": {0: batch, 1: seq_length},
                 "attention_mask": {0: batch, 2: "seq"},
-                "cache_position": {1: "seq"},
+                "cache_position": {0: "seq"},
                 "past_key_values": [
                     [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
                     [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],

onnx_diagnostic/torch_models/untrained/llm_tiny_llm.py

@@ -57,13 +57,13 @@ def get_tiny_llm(
     res = get_inputs(
         model,
         conf,
-        dummy_max_token_id=config["vocab_size"],
-        num_hidden_layers=config["num_hidden_layers"],
+        dummy_max_token_id=config["vocab_size"],  # type: ignore[arg-type]
+        num_hidden_layers=config["num_hidden_layers"],  # type: ignore[arg-type]
         batch_size=batch_size,
         sequence_length=sequence_length,
         sequence_length2=sequence_length2,
         dynamic_rope=dynamic_rope,
-        num_key_value_heads=config["num_key_value_heads"],
+        num_key_value_heads=config["num_key_value_heads"],  # type: ignore[arg-type]
         cls_cache="StaticCache" if use_static_cache else "DynamicCache",
     )