simple patch

Author: xadupre

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -426,14 +426,6 @@ def torch_export_patches(
                 patch_transformers_list, verbose=verbose
             )
 
-            if patch_transformers_list.patch_is_initialized:
-                if verbose:
-                    print(
-                        "[torch_export_patches] patches "
-                        "transformers.cache_utils.CacheLayerMixin.is_initialized"
-                    )
-                patch_transformers_list.apply_patch_for_is_initialized()
-
             if (
                 masking_utils
                 and patch_transformers_list.patch_masking_utils
@@ -697,14 +689,6 @@ def torch_export_patches(
                                 "in ALL_MASK_ATTENTION_FUNCTIONS"
                             )
 
-                if patch_transformers_list.patch_is_initialized:
-                    if verbose:
-                        print(
-                            "[torch_export_patches] restores "
-                            "transformers.cache_utils.CacheLayerMixin.is_initialized"
-                        )
-                    patch_transformers_list.disable_patch_for_is_initialized()
-
             ########
             # caches
             ########

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -38,35 +38,6 @@
 patch_is_initialized = pv.Version(transformers.__version__) > pv.Version("4.56.99")
 
 
-def _get_is_initialized(self):
-    return self.keys is not None
-
-
-def _set_is_initialized(self, value):
-    assert (value and self.keys is not None) or (not value and self.keys is None), (
-        f"The patch does not set is_initialized but checks the it is consistent with "
-        f"``self.keys is not None``, value={value}, "
-        f"self.keys is not None={self.keys is not None}"
-    )
-
-
-def apply_patch_for_is_initialized():
-    """
-    Fixes export issues introduced by PR `40791 <https://github.com/huggingface/transformers/pull/40791>`_.
-    The attribute is_initialized does not seem to be captured by :func:`torch.export.export`.
-    """
-    if patch_is_initialized:
-        transformers.cache_utils.CacheLayerMixin.is_initialized = property(
-            _get_is_initialized, _set_is_initialized
-        )
-
-
-def disable_patch_for_is_initialized():
-    """Disables the patch applied by function :func:`applies_patch_for_is_initialized`."""
-    if patch_is_initialized:
-        delattr(transformers.cache_utils.CacheLayerMixin, "is_initialized")
-
-
 if patch_masking_utils:
     # Introduced in 4.52
     from transformers.masking_utils import (
@@ -245,6 +216,8 @@ def lazy_initialization(self, key_states: torch.Tensor):
             new_shape[-2] = 0
             self.keys = torch.empty(new_shape, dtype=self.dtype, device=self.device)
             self.values = torch.empty(new_shape, dtype=self.dtype, device=self.device)
+            if patch_is_initialized:
+                self.is_initialized = True
 
 
 def _patch_make_causal_mask(