fix cachekeyvalue

Author: xadupre

_unittests/ut_helpers/test_cache_helper.py

@@ -50,6 +50,10 @@ def test_replace_by(self):
         past_key_values = make_dynamic_cache(
             [(torch.randn(bsize, nheads, slen, dim), torch.randn(bsize, nheads, slen, dim))]
         )
+        self.assertEqual(
+            "DynamicCache(key_cache=#1[T1s2x4x3x7], value_cache=#1[T1s2x4x3x7])",
+            self.string_type(past_key_values, with_shape=True),
+        )
         kwargs = dict(
             input_ids=torch.zeros(2, 3),
             attention_mask=torch.zeros(2, 3),

onnx_diagnostic/helpers/cache_helper.py

@@ -20,15 +20,27 @@ class CacheKeyValue:
     .. code-block:: python
 
         capi = CacheKeyValue(cache)
-        capi.cache.key_cache
-        capi.cache.value_cache
+        capi.key_cache
+        capi.value_cache
     """
 
     def __init__(self, cache):
         if hasattr(cache, "layers"):
-            layers = [layer for layer in cache.layers if layer is not None]
+            layers = [
+                layer
+                for layer in cache.layers
+                if layer is not None and layer.keys is not None and layer.values is not None
+            ]
             self.key_cache = [layer.keys for layer in layers]
             self.value_cache = [layer.values for layer in layers]
+            if None in self.key_cache or None in self.value_cache:
+                from .helper import string_type
+
+                raise AssertionError(
+                    f"issue with key_cache={string_type(self.key_cache)}, "
+                    f"or value_cache={string_type(self.value_cache)}, "
+                    f"cache.layers={string_type(cache.layers)}"
+                )
         else:
             self.key_cache = cache.key_cache
             self.value_cache = cache.value_cache

onnx_diagnostic/helpers/helper.py

@@ -710,6 +710,32 @@ def string_type(
         )
         return f"{obj.__class__.__name__}[{obj.cache_type}]{s}"
 
+    if obj.__class__.__name__ == "DynamicLayer":
+        import transformers
+
+        assert isinstance(
+            obj, transformers.cache_utils.DynamicLayer
+        ), f"Unexpected type {type(obj)}"
+        if verbose:
+            print(f"[string_type] LY0:{type(obj)}")
+        s1 = string_type(
+            obj.keys,
+            with_shape=with_shape,
+            with_min_max=with_min_max,
+            with_device=with_device,
+            limit=limit,
+            verbose=verbose,
+        )
+        s2 = string_type(
+            obj.values,
+            with_shape=with_shape,
+            with_min_max=with_min_max,
+            with_device=with_device,
+            limit=limit,
+            verbose=verbose,
+        )
+        return f"{obj.__class__.__name__}(keys={s1}, values={s2})"
+
     if isinstance(obj, torch.nn.Module):
         if verbose:
             print(f"[string_type] MM:{type(obj)}")