rewrite patch

Author: xadupre

onnx_diagnostic/helpers/cache_helper.py

@@ -28,14 +28,6 @@ def __init__(self, cache=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)}"
-                )
         elif cache is not None and hasattr(cache, "key_cache"):
             self.key_cache = cache.key_cache
             self.value_cache = cache.value_cache

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_dynamic_cache.py

@@ -27,20 +27,15 @@ def lazy_initialization(self, key_states: torch.Tensor):
             assert (
                 hasattr(key_states, "shape") and key_states is not None
             ), f"Attribute 'shape' is wrong for type {type(key_states)}"
-            assert isinstance(key_states.shape, tuple), (
-                f"Unxpected type {type(key_states.shape)} for key_states.shape, "
-                f"__dict__={key_states.shape.__dict__}"
-            )
-            new_shape = list(key_states.shape)
-            new_shape[-2] = 0
+            like = key_states[:, :0]
             # PATCHED: used a tensor with an empty shape and not en empty list to initialize
             if isinstance(key_states, torch._subclasses.fake_tensor.FakeTensor):
                 with key_states.fake_mode:
-                    self.keys = torch.empty(new_shape, dtype=self.dtype, device=self.device)
-                    self.values = torch.empty(new_shape, dtype=self.dtype, device=self.device)
+                    self.keys = torch.empty_like(like, dtype=self.dtype, device=self.device)
+                    self.values = torch.empty_like(like, dtype=self.dtype, device=self.device)
             else:
-                self.keys = torch.empty(new_shape, dtype=self.dtype, device=self.device)
-                self.values = torch.empty(new_shape, dtype=self.dtype, device=self.device)
+                self.keys = torch.empty_like(like, dtype=self.dtype, device=self.device)
+                self.values = torch.empty_like(like, dtype=self.dtype, device=self.device)
             if patch_is_initialized:
                 self.is_initialized = True