onnx-diagnostic 0.7.16__py3-none-any.whl → 0.8.1__py3-none-any.whl

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -1,5 +1,6 @@
 import functools
 import importlib
+import inspect
 import contextlib
 import re
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -8,6 +9,7 @@ from .onnx_export_serialization import (
     unregister_cache_serialization,
 )
 from .patches import patch_transformers as patch_transformers_list
+from .patch_details import PatchDetails
 
 
 def get_function(name: str) -> Tuple[type, Callable]:
@@ -51,7 +53,9 @@ def get_patches(mod, verbose: int = 0) -> Tuple[str, List[Any]]:
     return name, to_patch
 
 
-def patch_module_or_classes(mod, verbose: int = 0) -> Dict[type, Dict[type, Callable]]:
+def patch_module_or_classes(
+    mod, verbose: int = 0, patch_details: Optional[PatchDetails] = None
+) -> Dict[type, Dict[type, Callable]]:
     """
     Applies all patches defined in classes prefixed by ``patched_``
     ``cls._PATCHED_CLASS_`` defines the class to patch,
@@ -61,13 +65,16 @@ def patch_module_or_classes(mod, verbose: int = 0) -> Dict[type, Dict[type, Call
 
     :param mod: module of list of clsses to patch
     :param verbose: verbosity
+    :param patch_details: used to store information about the applied patches
     :return: patch info
     """
     if isinstance(mod, list):
         to_patch = mod
         name = "list"
+        list_name = "auto/list"
     else:
         name, to_patch = get_patches(mod, verbose)
+        list_name = f"auto/{mod.__name__.split('.')[-1]}"
 
     res = {}
     for cls in to_patch:
@@ -76,9 +83,15 @@ def patch_module_or_classes(mod, verbose: int = 0) -> Dict[type, Dict[type, Call
             keep = {}
             original = cls["module"]
             f = cls["function"]
+            assert not f.__name__.startswith("patched_"), (
+                f"The function {f} was already patched or the patch was not removed, "
+                f"original={original}"
+            )
             res[f] = f
             if verbose:
                 print(f"[patch_module_or_classes] function: {original.__name__}.{f.__name__}")
+            if patch_details:
+                patch_details.append(list_name, getattr(original, f.__name__), cls["patch"])
             setattr(original, f.__name__, cls["patch"])
             continue
 
@@ -89,6 +102,18 @@ def patch_module_or_classes(mod, verbose: int = 0) -> Dict[type, Dict[type, Call
 
         keep = {n: getattr(original, n, None) for n in methods}
         for n in methods:
+            if patch_details:
+                if hasattr(original, n):
+                    p = patch_details.append(list_name, getattr(original, n), getattr(cls, n))
+                else:
+                    p = patch_details.append(
+                        list_name, f"{original.__name__}{n}", getattr(cls, n)
+                    )
+                if "@patched_dynamic_rope_update" in inspect.getsource(getattr(cls, n)):
+                    # a tweak to include that patch.
+                    f = patch_details.find("patched_dynamic_rope_update")
+                    if f is not None:
+                        p.add_dependency(f)
             setattr(original, n, getattr(cls, n))
         res[cls] = keep
 
@@ -157,6 +182,628 @@ def register_additional_serialization_functions(
         unregister_cache_serialization(done, verbose=verbose)
 
 
+def _patch_sympy(verbose: int, patch_details: PatchDetails) -> Tuple[Optional[Callable], ...]:
+    import sympy
+
+    f_sympy_name = getattr(sympy.core.numbers.IntegerConstant, "name", None)
+
+    if verbose:
+        print(f"[torch_export_patches] sympy.__version__={sympy.__version__!r}")
+        print("[torch_export_patches] patch sympy")
+
+    sympy.core.numbers.IntegerConstant.name = lambda self: f"IntCst{str(self)}"
+    if patch_details:
+        patch_details.append(
+            "sympy",
+            f_sympy_name or "sympy.core.numbers.IntegerConstant.name",
+            sympy.core.numbers.IntegerConstant.name,
+        )
+    return (f_sympy_name,)
+
+
+def _unpatch_sympy(verbose: int, f_sympy_name: Optional[Callable]):
+    # tracked by https://github.com/pytorch/pytorch/issues/143494
+    import sympy
+
+    if f_sympy_name:
+        sympy.core.numbers.IntegerConstant.name = f_sympy_name
+    else:
+        delattr(sympy.core.numbers.IntegerConstant, "name")
+
+    if verbose:
+        print("[torch_export_patches] restored sympy functions")
+
+
+def _patch_torch(
+    verbose: int,
+    patch_details: PatchDetails,
+    patch_torch: int,
+    catch_constraints: bool,
+    stop_if_static: int,
+) -> Tuple[Optional[Callable], ...]:
+    import torch
+    import torch.jit
+    import torch._export.non_strict_utils  # produce_guards_and_solve_constraints
+    from torch.fx.experimental.symbolic_shapes import ShapeEnv
+    from .patches.patch_torch import (
+        patched_infer_size,
+        patched_vmap,
+        patched__broadcast_shapes,
+        patched__constrain_user_specified_dimhint_range,
+        _catch_produce_guards_and_solve_constraints,
+        patch__check_input_constraints_for_graph,
+        patched__broadcast_in_dim_meta,
+        patched__broadcast_in_dim_meta_level_2,
+        patched__maybe_broadcast,
+        patched_ShapeEnv,
+    )
+
+    f___constrain_user_specified_dimhint_range = None
+    f__broadcast_in_dim_meta = None
+    f__broadcast_shapes = None
+    f__check_input_constraints_for_graph = None
+    f__maybe_broadcast = None
+    f_broadcast_in_dim = None
+    f_infer_size = None
+    f_jit_isinstance = None
+    f_mark_static_address = None
+    f_produce_guards_and_solve_constraints = None
+    f_shape_env__check_frozen = None
+    f_shape_env__evaluate_expr = None
+    f_shape_env__log_guard = None
+    f_shape_env__set_replacement = None
+    f_vmap = None
+
+    if verbose:
+        print(f"[torch_export_patches] torch.__version__={torch.__version__!r}")
+        print(f"[torch_export_patches] stop_if_static={stop_if_static!r}")
+        print("[torch_export_patches] patch pytorch")
+
+    # torch.vmap
+    f_vmap = torch.vmap
+    torch.vmap = patched_vmap
+
+    # torch.jit.isinstance
+    f_jit_isinstance = torch.jit.isinstance
+    torch.jit.isinstance = isinstance
+
+    # torch._dynamo.mark_static_address
+    f_mark_static_address = torch._dynamo.mark_static_address
+    torch._dynamo.mark_static_address = lambda *_, **y_: None
+
+    # torch._subclasses.fake_impls.infer_size
+    f_infer_size = torch._subclasses.fake_impls.infer_size
+    torch._subclasses.fake_impls.infer_size = patched_infer_size
+    if patch_details:
+        patch_details.append("torch", f_infer_size, patched_infer_size)
+
+    # torch._refs._broadcast_shapes
+    f__broadcast_shapes = torch._refs._broadcast_shapes
+    torch._refs._broadcast_shapes = patched__broadcast_shapes
+    torch._meta_registrations._broadcast_shapes = patched__broadcast_shapes
+    if patch_details:
+        patch_details.append("torch", f__broadcast_shapes, patched__broadcast_shapes)
+
+    # torch._export.non_strict_utils._constrain_user_specified_dimhint_range
+    f___constrain_user_specified_dimhint_range = (
+        torch._export.non_strict_utils._constrain_user_specified_dimhint_range
+    )
+    torch._export.non_strict_utils._constrain_user_specified_dimhint_range = (
+        patched__constrain_user_specified_dimhint_range
+    )
+    if patch_details:
+        patch_details.append(
+            "torch",
+            f___constrain_user_specified_dimhint_range,
+            patched__constrain_user_specified_dimhint_range,
+        )
+
+    # torch._prims._broadcast_in_dim_meta
+    f_broadcast_in_dim = torch._prims.broadcast_in_dim
+    f__broadcast_in_dim_meta = torch._prims._broadcast_in_dim_meta
+    _patched_dim_f = (
+        patched__broadcast_in_dim_meta_level_2
+        if patch_torch == 2
+        else patched__broadcast_in_dim_meta
+    )
+    torch._prims._broadcast_in_dim_meta = _patched_dim_f
+    torch._prims.broadcast_in_dim = _patched_dim_f
+    if patch_details:
+        patch_details.append("torch", f__broadcast_in_dim_meta, _patched_dim_f)
+
+    # torch._refs._maybe_broadcast
+    f__maybe_broadcast = torch._refs._maybe_broadcast
+    torch._refs._maybe_broadcast = patched__maybe_broadcast
+    if patch_details:
+        patch_details.append("torch", f__maybe_broadcast, patched__maybe_broadcast)
+
+    # ShapeEnv
+    f_shape_env__evaluate_expr = ShapeEnv._evaluate_expr
+    ShapeEnv._evaluate_expr = patched_ShapeEnv._evaluate_expr
+    if patch_details:
+        patch_details.append(
+            "torch", f_shape_env__evaluate_expr, patched_ShapeEnv._evaluate_expr
+        )
+
+    # torch._export.non_strict_utils.produce_guards_and_solve_constraints
+    if catch_constraints:
+        if verbose:
+            print("[torch_export_patches] modifies shape constraints")
+        f_produce_guards_and_solve_constraints = (
+            torch._export.non_strict_utils.produce_guards_and_solve_constraints
+        )
+        f__check_input_constraints_for_graph = (
+            torch._export.utils._check_input_constraints_for_graph
+        )
+        torch._export.non_strict_utils.produce_guards_and_solve_constraints = (
+            lambda *args, **kwargs: _catch_produce_guards_and_solve_constraints(
+                f_produce_guards_and_solve_constraints, *args, verbose=verbose, **kwargs
+            )
+        )
+        torch._export.utils._check_input_constraints_for_graph = (
+            lambda *args, **kwargs: patch__check_input_constraints_for_graph(
+                f__check_input_constraints_for_graph, *args, verbose=verbose, **kwargs
+            )
+        )
+
+    if patch_torch and stop_if_static:
+        ShapeEnv._log_guard_remember = ShapeEnv._log_guard
+
+        if verbose:
+            print("[torch_export_patches] assert when a dynamic dimension turns static")
+            print("[torch_export_patches] replaces ShapeEnv._set_replacement")
+
+        f_shape_env__set_replacement = ShapeEnv._set_replacement
+        ShapeEnv._set_replacement = patched_ShapeEnv._set_replacement
+        if patch_details:
+            patch_details.append(
+                "torch", f_shape_env__set_replacement, patched_ShapeEnv._set_replacement
+            )
+
+        if verbose:
+            print("[torch_export_patches] replaces ShapeEnv._log_guard")
+        f_shape_env__log_guard = ShapeEnv._log_guard
+        ShapeEnv._log_guard = patched_ShapeEnv._log_guard
+        if patch_details:
+            patch_details.append("torch", f_shape_env__log_guard, patched_ShapeEnv._log_guard)
+
+        if stop_if_static > 1:
+            if verbose:
+                print("[torch_export_patches] replaces ShapeEnv._check_frozen")
+            f_shape_env__check_frozen = ShapeEnv._check_frozen
+            ShapeEnv._check_frozen = patched_ShapeEnv._check_frozen
+            if patch_details:
+                patch_details.append(
+                    "torch", f_shape_env__check_frozen, ShapeEnv._check_frozen
+                )
+    return (
+        f___constrain_user_specified_dimhint_range,
+        f__broadcast_in_dim_meta,
+        f__broadcast_shapes,
+        f__check_input_constraints_for_graph,
+        f__maybe_broadcast,
+        f_broadcast_in_dim,
+        f_infer_size,
+        f_jit_isinstance,
+        f_mark_static_address,
+        f_produce_guards_and_solve_constraints,
+        f_shape_env__check_frozen,
+        f_shape_env__evaluate_expr,
+        f_shape_env__log_guard,
+        f_shape_env__set_replacement,
+        f_vmap,
+    )
+
+
+def _unpatch_torch(
+    verbose: int,
+    _patch_details: PatchDetails,
+    patch_torch: int,
+    catch_constraints: bool,
+    stop_if_static: int,
+    f___constrain_user_specified_dimhint_range: Optional[Callable],
+    f__broadcast_in_dim_meta: Optional[Callable],
+    f__broadcast_shapes: Optional[Callable],
+    f__check_input_constraints_for_graph: Optional[Callable],
+    f__maybe_broadcast: Optional[Callable],
+    f_broadcast_in_dim: Optional[Callable],
+    f_infer_size: Optional[Callable],
+    f_jit_isinstance: Optional[Callable],
+    f_mark_static_address: Optional[Callable],
+    f_produce_guards_and_solve_constraints: Optional[Callable],
+    f_shape_env__check_frozen: Optional[Callable],
+    f_shape_env__evaluate_expr: Optional[Callable],
+    f_shape_env__log_guard: Optional[Callable],
+    f_shape_env__set_replacement: Optional[Callable],
+    f_vmap: Optional[Callable],
+):
+    import torch
+    import torch.jit
+    import torch._export.non_strict_utils  # produce_guards_and_solve_constraints
+    from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+    # this should disappear when torch.jit is removed
+    torch.vmap = f_vmap
+    torch.jit.isinstance = f_jit_isinstance
+    torch._dynamo.mark_static_address = f_mark_static_address
+    # tracked by https://github.com/pytorch/pytorch/issues/143495
+    torch._subclasses.fake_impls.infer_size = f_infer_size
+    torch._refs._broadcast_shapes = f__broadcast_shapes
+    torch._meta_registrations._broadcast_shapes = f__broadcast_shapes
+    torch._export.non_strict_utils._constrain_user_specified_dimhint_range = (
+        f___constrain_user_specified_dimhint_range
+    )
+    torch._prims._broadcast_in_dim_meta = f__broadcast_in_dim_meta
+    torch._prims.broadcast_in_dim = f_broadcast_in_dim
+    torch._refs._maybe_broadcast = f__maybe_broadcast
+    ShapeEnv._evaluate_expr = f_shape_env__evaluate_expr
+
+    if verbose:
+        print("[torch_export_patches] restored pytorch functions")
+
+    if patch_torch and stop_if_static:
+        if verbose:
+            print("[torch_export_patches] restored ShapeEnv._set_replacement")
+
+        ShapeEnv._set_replacement = f_shape_env__set_replacement
+
+        if verbose:
+            print("[torch_export_patches] restored ShapeEnv._log_guard")
+
+        ShapeEnv._log_guard = f_shape_env__log_guard
+
+        if stop_if_static > 1:
+            if verbose:
+                print("[torch_export_patches] restored ShapeEnv._check_frozen")
+            ShapeEnv._check_frozen = f_shape_env__check_frozen
+
+    if patch_torch and catch_constraints:
+        # to catch or skip dynamic_shapes issues
+        torch._export.non_strict_utils.produce_guards_and_solve_constraints = (
+            f_produce_guards_and_solve_constraints
+        )
+        torch._export.utils._check_input_constraints_for_graph = (
+            f__check_input_constraints_for_graph
+        )
+        if verbose:
+            print("[torch_export_patches] restored shape constraints")
+
+
+def _patch_transformers(
+    verbose: int, patch_details: PatchDetails
+) -> Tuple[Optional[Callable], ...]:
+    import transformers
+
+    try:
+        import transformers.masking_utils as masking_utils
+    except ImportError:
+        masking_utils = None
+
+    try:
+        import transformers.integrations.sdpa_attention as sdpa_attention
+    except ImportError:
+        sdpa_attention = None
+
+    try:
+        import transformers.modeling_utils as modeling_utils
+    except ImportError:
+        modeling_utils = None
+
+    try:
+        import transformers.modeling_rope_utils as modeling_rope_utils
+    except ImportError:
+        modeling_rope_utils = None
+
+    if (
+        patch_details
+        and modeling_rope_utils
+        and hasattr(modeling_rope_utils, "dynamic_rope_update")
+    ):
+        patch_details.append(
+            "patch_transformers",
+            modeling_rope_utils.dynamic_rope_update,
+            patch_transformers_list.patched_dynamic_rope_update,
+        )
+
+    if verbose:
+        print(f"[torch_export_patches] transformers.__version__={transformers.__version__!r}")
+    assert not sdpa_attention.sdpa_attention_forward.__name__.startswith("patched_"), (
+        f"Function 'sdpa_attention.sdpa_attention_forward' is already patched, "
+        f"sdpa_attention.sdpa_attention_forward={sdpa_attention.sdpa_attention_forward}"
+    )
+
+    f_transformers__vmap_for_bhqkv = None
+    f_transformers_eager_mask = None
+    f_transformers_sdpa_attention_forward = None
+    f_transformers_sdpa_mask = None
+    f_transformers_sdpa_mask_recent_torch = None
+
+    if (  # vmap
+        masking_utils
+        and patch_transformers_list.patch_masking_utils
+        and hasattr(masking_utils, "_vmap_for_bhqkv")
+    ):
+        if verbose:
+            print("[torch_export_patches] patches transformers.masking_utils._vmap_for_bhqkv")
+        f_transformers__vmap_for_bhqkv = masking_utils._vmap_for_bhqkv
+        masking_utils._vmap_for_bhqkv = patch_transformers_list.patched__vmap_for_bhqkv
+        if patch_details:
+            patch_details.append(
+                "transformers",
+                f_transformers__vmap_for_bhqkv,
+                patch_transformers_list.patched__vmap_for_bhqkv,
+            )
+
+        if verbose:
+            print(
+                "[torch_export_patches] patches "
+                "transformers.masking_utils.sdpa_mask_recent_torch"
+            )
+        f_transformers_sdpa_mask_recent_torch = masking_utils.sdpa_mask_recent_torch
+        masking_utils.sdpa_mask_recent_torch = (
+            patch_transformers_list.patched_sdpa_mask_recent_torch
+        )
+        if patch_details:
+            patch_details.append(
+                "transformers",
+                f_transformers_sdpa_mask_recent_torch,
+                patch_transformers_list.patched_sdpa_mask_recent_torch,
+            )
+        if masking_utils.sdpa_mask == f_transformers_sdpa_mask_recent_torch:
+            if verbose:
+                print("[torch_export_patches] patches transformers.masking_utils.sdpa_mask")
+            f_transformers_sdpa_mask = masking_utils.sdpa_mask
+            masking_utils.sdpa_mask = patch_transformers_list.patched_sdpa_mask_recent_torch
+            if patch_details:
+                patch_details.append(
+                    "transformers",
+                    f_transformers_sdpa_mask,
+                    patch_transformers_list.patched_sdpa_mask_recent_torch,
+                )
+        else:
+            f_transformers_sdpa_mask = None
+
+    if (  # eager_mask
+        masking_utils
+        and patch_transformers_list.patch_masking_utils
+        and hasattr(masking_utils, "eager_mask")
+    ):
+        if verbose:
+            print("[torch_export_patches] patches transformers.masking_utils.eager_mask")
+        f_transformers_eager_mask = masking_utils.eager_mask
+        masking_utils.eager_mask = patch_transformers_list.patched_eager_mask
+        if patch_details:
+            patch_details.append(
+                "transformers",
+                f_transformers_eager_mask,
+                patch_transformers_list.patched_eager_mask,
+            )
+        if (
+            "eager" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS
+            and masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["eager"]
+            == f_transformers_eager_mask
+        ):
+            if verbose:
+                print(
+                    "[torch_export_patches] patches "
+                    "transformers.masking_utils.eager_mask "
+                    "in ALL_MASK_ATTENTION_FUNCTIONS"
+                )
+            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["eager"] = (
+                patch_transformers_list.patched_eager_mask
+            )
+
+    if (  # sdpa_mask
+        masking_utils
+        and patch_transformers_list.patch_masking_utils
+        and hasattr(masking_utils, "sdpa_mask")
+        and f_transformers_sdpa_mask is not None
+    ):
+        if verbose:
+            print(
+                "[torch_export_patches] patches "
+                "transformers.masking_utils.sdpa_mask "
+                "in ALL_MASK_ATTENTION_FUNCTIONS"
+            )
+        if (
+            "sdpa" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS
+            and masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] == f_transformers_sdpa_mask
+        ):
+            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] = (
+                patch_transformers_list.patched_sdpa_mask_recent_torch
+            )
+
+    if (  # sdpa_attention_forward
+        sdpa_attention is not None
+        and modeling_utils is not None
+        and hasattr(sdpa_attention, "sdpa_attention_forward")
+        and hasattr(sdpa_attention, "use_gqa_in_sdpa")
+        and hasattr(modeling_utils, "AttentionInterface")
+    ):
+        if verbose:
+            print(
+                "[torch_export_patches] patches "
+                "transformers.integrations.sdpa_attention.sdpa_attention_forward"
+            )
+        f_transformers_sdpa_attention_forward = sdpa_attention.sdpa_attention_forward
+        assert not f_transformers_sdpa_attention_forward.__name__.startswith("patched_"), (
+            f"Function 'sdpa_attention.sdpa_attention_forward' is already patched, "
+            f"sdpa_attention.sdpa_attention_forward={f_transformers_sdpa_attention_forward}"
+        )
+        sdpa_attention.sdpa_attention_forward = (
+            patch_transformers_list.patched_sdpa_attention_forward
+        )
+        modeling_utils.sdpa_attention_forward = (
+            patch_transformers_list.patched_sdpa_attention_forward
+        )
+        modeling_utils.AttentionInterface._global_mapping["sdpa"] = (
+            patch_transformers_list.patched_sdpa_attention_forward
+        )
+        if patch_details:
+            patch_details.append(
+                "transformers",
+                f_transformers_sdpa_attention_forward,
+                patch_transformers_list.patched_sdpa_attention_forward,
+            )
+
+    revert_patches_info = patch_module_or_classes(
+        patch_transformers_list, verbose=verbose, patch_details=patch_details
+    )
+
+    return (
+        f_transformers__vmap_for_bhqkv,
+        f_transformers_eager_mask,
+        f_transformers_sdpa_attention_forward,
+        f_transformers_sdpa_mask,
+        f_transformers_sdpa_mask_recent_torch,
+        revert_patches_info,
+    )
+
+
+def _unpatch_transformers(
+    verbose: int,
+    _patch_details: PatchDetails,
+    f_transformers__vmap_for_bhqkv: Optional[Callable],
+    f_transformers_eager_mask: Optional[Callable],
+    f_transformers_sdpa_attention_forward: Optional[Callable],
+    f_transformers_sdpa_mask: Optional[Callable],
+    f_transformers_sdpa_mask_recent_torch: Optional[Callable],
+    revert_patches_info: Optional[Callable],
+):
+
+    try:
+        import transformers.masking_utils as masking_utils
+    except ImportError:
+        masking_utils = None
+
+    try:
+        import transformers.integrations.sdpa_attention as sdpa_attention
+    except ImportError:
+        sdpa_attention = None
+
+    try:
+        import transformers.modeling_utils as modeling_utils
+    except ImportError:
+        modeling_utils = None
+
+    try:
+        import transformers.masking_utils as masking_utils
+    except ImportError:
+        masking_utils = None
+    if verbose:
+        print("[torch_export_patches] unpatches transformers")
+
+    if (  # vmap
+        masking_utils
+        and patch_transformers_list.patch_masking_utils
+        and hasattr(masking_utils, "_vmap_for_bhqkv")
+    ):
+        assert f_transformers__vmap_for_bhqkv.__name__ == "_vmap_for_bhqkv", (
+            f"corrupted function '_vmap_for_bhqkv', its name is "
+            f"{f_transformers__vmap_for_bhqkv.__name__!r}"
+        )
+        masking_utils._vmap_for_bhqkv = f_transformers__vmap_for_bhqkv
+
+        if verbose:
+            print("[torch_export_patches] restored transformers.masking_utils._vmap_for_bhqkv")
+
+        assert f_transformers_sdpa_mask_recent_torch.__name__ == "sdpa_mask_recent_torch", (
+            f"corrupted function 'sdpa_mask_recent_torch', its name is "
+            f"{f_transformers_sdpa_mask_recent_torch.__name__!r}"
+        )
+        masking_utils.sdpa_mask_recent_torch = f_transformers_sdpa_mask_recent_torch
+
+        if verbose:
+            print(
+                "[torch_export_patches] restored "
+                "transformers.masking_utils.sdpa_mask_recent_torch"
+            )
+
+        if f_transformers_sdpa_mask is not None:
+            assert f_transformers_sdpa_mask.__name__ in (
+                "sdpa_mask",
+                "sdpa_mask_recent_torch",
+            ), (
+                f"corrupted function 'sdpa_mask', its name is "
+                f"{f_transformers_sdpa_mask.__name__!r}"
+            )
+            masking_utils.sdpa_mask = f_transformers_sdpa_mask
+            if verbose:
+                print("[torch_export_patches] restored transformers.masking_utils.sdpa_mask")
+
+    if (  # eager_mask
+        masking_utils
+        and patch_transformers_list.patch_masking_utils
+        and hasattr(masking_utils, "eager_mask")
+    ):
+        assert f_transformers_eager_mask.__name__ == "eager_mask", (
+            f"corrupted function 'eager_mask', its name is "
+            f"{f_transformers_eager_mask.__name__!r}"
+        )
+        masking_utils.eager_mask = f_transformers_eager_mask
+        if verbose:
+            print("[torch_export_patches] restored transformers.masking_utils.eager_mask")
+        assert masking_utils.eager_mask.__name__ == "eager_mask", (
+            f"corrupted function 'eager_mask', its name is "
+            f"{masking_utils.eager_mask.__name__!r}"
+        )
+        if (
+            "eager" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS
+            and masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["eager"]
+            == patch_transformers_list.patched_eager_mask
+        ):
+            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["eager"] = f_transformers_eager_mask
+            if verbose:
+                print(
+                    "[torch_export_patches] restored "
+                    "transformers.masking_utils.eager_mask "
+                    "in ALL_MASK_ATTENTION_FUNCTIONS"
+                )
+        assert masking_utils.eager_mask.__name__ == "eager_mask", (
+            f"corrupted function 'eager_mask', its name is "
+            f"{masking_utils.eager_mask.__name__!r}"
+        )
+
+    if (  # sdpa_mask
+        masking_utils
+        and patch_transformers_list.patch_masking_utils
+        and hasattr(masking_utils, "sdpa_mask")
+    ):
+        if (
+            "sdpa" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS
+            and masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"]
+            == patch_transformers_list.patched_sdpa_mask_recent_torch
+        ):
+            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] = f_transformers_sdpa_mask
+            if verbose:
+                print(
+                    "[torch_export_patches] restored "
+                    "transformers.masking_utils.sdpa_mask "
+                    "in ALL_MASK_ATTENTION_FUNCTIONS"
+                )
+
+    if (  # sdpa_attention_forward
+        sdpa_attention is not None
+        and modeling_utils is not None
+        and hasattr(sdpa_attention, "sdpa_attention_forward")
+        and hasattr(sdpa_attention, "use_gqa_in_sdpa")
+        and hasattr(modeling_utils, "AttentionInterface")
+    ):
+        sdpa_attention.sdpa_attention_forward = f_transformers_sdpa_attention_forward
+        modeling_utils.sdpa_attention_forward = f_transformers_sdpa_attention_forward
+        modeling_utils.AttentionInterface._global_mapping["sdpa"] = (
+            f_transformers_sdpa_attention_forward
+        )
+        if verbose:
+            print(
+                "[torch_export_patches] restored "
+                "transformers.integrations.sdpa_attention."
+                "sdpa_attention_forward"
+            )
+
+    unpatch_module_or_classes(patch_transformers_list, revert_patches_info, verbose=verbose)
+
+
 @contextlib.contextmanager
 def torch_export_patches(
     patch_sympy: bool = True,
@@ -170,6 +817,7 @@ def torch_export_patches(
     custom_patches: Optional[List[type["torch.nn.Module"]]] = None,  # noqa: F821
     rewrite: Optional[List[Callable]] = None,
     dump_rewriting: Optional[str] = None,
+    patch_details: Optional[PatchDetails] = None,
 ) -> Callable:
     """
     Tries to bypass some situations :func:`torch.export.export` does not support.
@@ -200,6 +848,7 @@ def torch_export_patches(
         <onnx_diagnostic.torch_export_patches.patch_module.transform_method>`,
         its documentation provides possible values
     :param dump_rewriting: dumps rewriting information in file beginning with that prefix
+    :param patch_details: if specified, this class is used to stored every rewritten done.
     :param verbose: to show which patches is applied
 
     The list of available patches.
@@ -270,7 +919,10 @@ def torch_export_patches(
 
         with (
             torch_export_rewrite(
-                rewrite=rewrite, dump_rewriting=dump_rewriting, verbose=verbose
+                rewrite=rewrite,
+                dump_rewriting=dump_rewriting,
+                verbose=verbose,
+                patch_details=patch_details,
             ),
             torch_export_patches(  # type: ignore[var-annotated]
                 patch_sympy=patch_sympy,
@@ -282,6 +934,7 @@ def torch_export_patches(
                 verbose=verbose,
                 patch=patch,
                 custom_patches=custom_patches,
+                patch_details=patch_details,
             ) as f,
         ):
             try:
@@ -300,19 +953,13 @@ def torch_export_patches(
         finally:
             unregister_cache_serialization(done, verbose=verbose)
     else:
-        import torch
-        import torch._export.non_strict_utils  # produce_guards_and_solve_constraints
-        import torch.jit
-
         if verbose:
             print(
                 "[torch_export_patches] replace torch.jit.isinstance, "
                 "torch._dynamo.mark_static_address"
             )
 
-        ########
         # caches
-        ########
 
         cache_done = register_cache_serialization(
             patch_transformers=patch_transformers,
@@ -320,282 +967,50 @@ def torch_export_patches(
             verbose=verbose,
         )
 
-        #############
-        # patch sympy
-        #############
+        # patches
 
         if patch_sympy:
-            import sympy
-
-            f_sympy_name = getattr(sympy.core.numbers.IntegerConstant, "name", None)
-
-            if verbose:
-                print(f"[torch_export_patches] sympy.__version__={sympy.__version__!r}")
-                print("[torch_export_patches] patch sympy")
-
-            sympy.core.numbers.IntegerConstant.name = lambda self: f"IntCst{str(self)}"
-
-        ###############
-        # patch pytorch
-        ###############
+            (f_sympy_name,) = _patch_sympy(verbose, patch_details)
 
         if patch_torch:
-            from torch.fx.experimental.symbolic_shapes import ShapeEnv
-            from .patches.patch_torch import (
-                patched_infer_size,
-                patched_vmap,
-                patched__broadcast_shapes,
-                patched__constrain_user_specified_dimhint_range,
-                _catch_produce_guards_and_solve_constraints,
-                patch__check_input_constraints_for_graph,
-                patched__broadcast_in_dim_meta,
-                patched__broadcast_in_dim_meta_level_2,
-                patched__maybe_broadcast,
-                patched_ShapeEnv,
-            )
-
-            if verbose:
-                print(f"[torch_export_patches] torch.__version__={torch.__version__!r}")
-                print(f"[torch_export_patches] stop_if_static={stop_if_static!r}")
-                print("[torch_export_patches] patch pytorch")
-
-            # torch.vmap
-            f_vmap = torch.vmap
-            torch.vmap = patched_vmap
-
-            # torch.jit.isinstance
-            f_jit_isinstance = torch.jit.isinstance
-            torch.jit.isinstance = isinstance
-
-            # torch._dynamo.mark_static_address
-            f_mark_static_address = torch._dynamo.mark_static_address
-            torch._dynamo.mark_static_address = lambda *_, **y_: None
-
-            # torch._subclasses.fake_impls.infer_size
-            f_infer_size = torch._subclasses.fake_impls.infer_size
-            torch._subclasses.fake_impls.infer_size = patched_infer_size
-
-            # torch._refs._broadcast_shapes
-            f__broadcast_shapes = torch._refs._broadcast_shapes
-            torch._refs._broadcast_shapes = patched__broadcast_shapes
-            torch._meta_registrations._broadcast_shapes = patched__broadcast_shapes
-
-            # torch._export.non_strict_utils._constrain_user_specified_dimhint_range
-            f___constrain_user_specified_dimhint_range = (
-                torch._export.non_strict_utils._constrain_user_specified_dimhint_range
-            )
-            torch._export.non_strict_utils._constrain_user_specified_dimhint_range = (
-                patched__constrain_user_specified_dimhint_range
-            )
-
-            # torch._prims._broadcast_in_dim_meta
-            f_broadcast_in_dim = torch._prims.broadcast_in_dim
-            f__broadcast_in_dim_meta = torch._prims._broadcast_in_dim_meta
-            _patched_dim_f = (
-                patched__broadcast_in_dim_meta_level_2
-                if patch_torch == 2
-                else patched__broadcast_in_dim_meta
+            (
+                f___constrain_user_specified_dimhint_range,
+                f__broadcast_in_dim_meta,
+                f__broadcast_shapes,
+                f__check_input_constraints_for_graph,
+                f__maybe_broadcast,
+                f_broadcast_in_dim,
+                f_infer_size,
+                f_jit_isinstance,
+                f_mark_static_address,
+                f_produce_guards_and_solve_constraints,
+                f_shape_env__check_frozen,
+                f_shape_env__evaluate_expr,
+                f_shape_env__log_guard,
+                f_shape_env__set_replacement,
+                f_vmap,
+            ) = _patch_torch(
+                verbose, patch_details, patch_torch, catch_constraints, stop_if_static
             )
-            torch._prims._broadcast_in_dim_meta = _patched_dim_f
-            torch._prims.broadcast_in_dim = _patched_dim_f
-
-            # torch._refs._maybe_broadcast
-            f__maybe_broadcast = torch._refs._maybe_broadcast
-            torch._refs._maybe_broadcast = patched__maybe_broadcast
-
-            # ShapeEnv
-            f_shape_env__evaluate_expr = ShapeEnv._evaluate_expr
-            ShapeEnv._evaluate_expr = patched_ShapeEnv._evaluate_expr
-
-        # torch._export.non_strict_utils.produce_guards_and_solve_constraints
-        if patch_torch and catch_constraints:
-            if verbose:
-                print("[torch_export_patches] modifies shape constraints")
-            f_produce_guards_and_solve_constraints = (
-                torch._export.non_strict_utils.produce_guards_and_solve_constraints
-            )
-            f__check_input_constraints_for_graph = (
-                torch._export.utils._check_input_constraints_for_graph
-            )
-            torch._export.non_strict_utils.produce_guards_and_solve_constraints = (
-                lambda *args, **kwargs: _catch_produce_guards_and_solve_constraints(
-                    f_produce_guards_and_solve_constraints, *args, verbose=verbose, **kwargs
-                )
-            )
-            torch._export.utils._check_input_constraints_for_graph = (
-                lambda *args, **kwargs: patch__check_input_constraints_for_graph(
-                    f__check_input_constraints_for_graph, *args, verbose=verbose, **kwargs
-                )
-            )
-
-        if patch_torch and stop_if_static:
-            ShapeEnv._log_guard_remember = ShapeEnv._log_guard
-
-            if verbose:
-                print("[torch_export_patches] assert when a dynamic dimension turns static")
-                print("[torch_export_patches] replaces ShapeEnv._set_replacement")
-
-            f_shape_env__set_replacement = ShapeEnv._set_replacement
-            ShapeEnv._set_replacement = patched_ShapeEnv._set_replacement
-
-            if verbose:
-                print("[torch_export_patches] replaces ShapeEnv._log_guard")
-            f_shape_env__log_guard = ShapeEnv._log_guard
-            ShapeEnv._log_guard = patched_ShapeEnv._log_guard
-
-            if stop_if_static > 1:
-                if verbose:
-                    print("[torch_export_patches] replaces ShapeEnv._check_frozen")
-                f_shape_env__check_frozen = ShapeEnv._check_frozen
-                ShapeEnv._check_frozen = patched_ShapeEnv._check_frozen
-
-        ####################
-        # patch transformers
-        ####################
 
         if patch_transformers:
-            try:
-                import transformers.masking_utils as masking_utils
-            except ImportError:
-                masking_utils = None
-
-            try:
-                import transformers.integrations.sdpa_attention as sdpa_attention
-            except ImportError:
-                sdpa_attention = None
-
-            try:
-                import transformers.modeling_utils as modeling_utils
-            except ImportError:
-                modeling_utils = None
-
-            if verbose:
-                import transformers
-
-                print(
-                    f"[torch_export_patches] transformers.__version__="
-                    f"{transformers.__version__!r}"
-                )
-            revert_patches_info = patch_module_or_classes(
-                patch_transformers_list, verbose=verbose
-            )
-
-            if (  # vmap
-                masking_utils
-                and patch_transformers_list.patch_masking_utils
-                and hasattr(masking_utils, "_vmap_for_bhqkv")
-            ):
-                if verbose:
-                    print(
-                        "[torch_export_patches] patches "
-                        "transformers.masking_utils._vmap_for_bhqkv"
-                    )
-                f_transformers__vmap_for_bhqkv = masking_utils._vmap_for_bhqkv
-                masking_utils._vmap_for_bhqkv = patch_transformers_list.patched__vmap_for_bhqkv
-
-                if verbose:
-                    print(
-                        "[torch_export_patches] patches "
-                        "transformers.masking_utils.sdpa_mask_recent_torch"
-                    )
-                f_transformers_sdpa_mask_recent_torch = masking_utils.sdpa_mask_recent_torch
-                masking_utils.sdpa_mask_recent_torch = (
-                    patch_transformers_list.patched_sdpa_mask_recent_torch
-                )
-                if masking_utils.sdpa_mask == f_transformers_sdpa_mask_recent_torch:
-                    if verbose:
-                        print(
-                            "[torch_export_patches] patches "
-                            "transformers.masking_utils.sdpa_mask"
-                        )
-                    f_transformers_sdpa_mask = masking_utils.sdpa_mask
-                    masking_utils.sdpa_mask = (
-                        patch_transformers_list.patched_sdpa_mask_recent_torch
-                    )
-                else:
-                    f_transformers_sdpa_mask = None
-
-            if (  # eager_mask
-                masking_utils
-                and patch_transformers_list.patch_masking_utils
-                and hasattr(masking_utils, "eager_mask")
-            ):
-                if verbose:
-                    print(
-                        "[torch_export_patches] patches "
-                        "transformers.masking_utils.eager_mask"
-                    )
-                f_transformers_eager_mask = masking_utils.eager_mask
-                masking_utils.eager_mask = patch_transformers_list.patched_eager_mask
-                if (
-                    "eager" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS
-                    and masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["eager"]
-                    == f_transformers_eager_mask
-                ):
-                    if verbose:
-                        print(
-                            "[torch_export_patches] patches "
-                            "transformers.masking_utils.eager_mask "
-                            "in ALL_MASK_ATTENTION_FUNCTIONS"
-                        )
-                    masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["eager"] = (
-                        patch_transformers_list.patched_eager_mask
-                    )
-
-            if (  # sdpa_mask
-                masking_utils
-                and patch_transformers_list.patch_masking_utils
-                and hasattr(masking_utils, "sdpa_mask")
-                and f_transformers_sdpa_mask is not None
-            ):
-                if verbose:
-                    print(
-                        "[torch_export_patches] patches "
-                        "transformers.masking_utils.sdpa_mask "
-                        "in ALL_MASK_ATTENTION_FUNCTIONS"
-                    )
-                if (
-                    "sdpa" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS
-                    and masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"]
-                    == f_transformers_sdpa_mask
-                ):
-                    masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] = (
-                        patch_transformers_list.patched_sdpa_mask_recent_torch
-                    )
-
-            if (  # sdpa_attention_forward
-                sdpa_attention is not None
-                and modeling_utils is not None
-                and hasattr(sdpa_attention, "sdpa_attention_forward")
-                and hasattr(sdpa_attention, "use_gqa_in_sdpa")
-                and hasattr(modeling_utils, "AttentionInterface")
-            ):
-                if verbose:
-                    print(
-                        "[torch_export_patches] patches "
-                        "transformers.integrations.sdpa_attention.sdpa_attention_forward"
-                    )
-                f_sdpa_attention_forward = sdpa_attention.sdpa_attention_forward
-                sdpa_attention.sdpa_attention_forward = (
-                    patch_transformers_list.patched_sdpa_attention_forward
-                )
-                modeling_utils.sdpa_attention_forward = (
-                    patch_transformers_list.patched_sdpa_attention_forward
-                )
-                modeling_utils.AttentionInterface._global_mapping["sdpa"] = (
-                    patch_transformers_list.patched_sdpa_attention_forward
-                )
+            (
+                f_transformers__vmap_for_bhqkv,
+                f_transformers_eager_mask,
+                f_transformers_sdpa_attention_forward,
+                f_transformers_sdpa_mask,
+                f_transformers_sdpa_mask_recent_torch,
+                revert_patches_info,
+            ) = _patch_transformers(verbose, patch_details)
 
         if custom_patches:
             if verbose:
                 print("[torch_export_patches] applies custom patches")
             revert_custom_patches_info = patch_module_or_classes(
-                custom_patches, verbose=verbose
+                custom_patches, verbose=verbose, patch_details=patch_details
             )
 
-        ########
         # export
-        ########
 
         fct_callable = replacement_before_exporting if patch_transformers else (lambda x: x)
 
@@ -605,73 +1020,50 @@ def torch_export_patches(
         try:
             yield fct_callable
         finally:
-            #######
-            # sympy
-            #######
+
+            # unpatch
 
             if verbose:
                 print("[torch_export_patches] remove patches")
 
             if patch_sympy:
-                # tracked by https://github.com/pytorch/pytorch/issues/143494
-                if f_sympy_name:
-                    sympy.core.numbers.IntegerConstant.name = f_sympy_name
-                else:
-                    delattr(sympy.core.numbers.IntegerConstant, "name")
-
-                if verbose:
-                    print("[torch_export_patches] restored sympy functions")
-
-            #######
-            # torch
-            #######
+                _unpatch_sympy(verbose, f_sympy_name)
 
             if patch_torch:
-                # this should disappear when torch.jit is removed
-                torch.vmap = f_vmap
-                torch.jit.isinstance = f_jit_isinstance
-                torch._dynamo.mark_static_address = f_mark_static_address
-                # tracked by https://github.com/pytorch/pytorch/issues/143495
-                torch._subclasses.fake_impls.infer_size = f_infer_size
-                torch._refs._broadcast_shapes = f__broadcast_shapes
-                torch._meta_registrations._broadcast_shapes = f__broadcast_shapes
-                torch._export.non_strict_utils._constrain_user_specified_dimhint_range = (
-                    f___constrain_user_specified_dimhint_range
+                _unpatch_torch(
+                    verbose,
+                    patch_details,
+                    patch_torch,
+                    catch_constraints,
+                    stop_if_static,
+                    f___constrain_user_specified_dimhint_range,
+                    f__broadcast_in_dim_meta,
+                    f__broadcast_shapes,
+                    f__check_input_constraints_for_graph,
+                    f__maybe_broadcast,
+                    f_broadcast_in_dim,
+                    f_infer_size,
+                    f_jit_isinstance,
+                    f_mark_static_address,
+                    f_produce_guards_and_solve_constraints,
+                    f_shape_env__check_frozen,
+                    f_shape_env__evaluate_expr,
+                    f_shape_env__log_guard,
+                    f_shape_env__set_replacement,
+                    f_vmap,
                 )
-                torch._prims._broadcast_in_dim_meta = f__broadcast_in_dim_meta
-                torch._prims.broadcast_in_dim = f_broadcast_in_dim
-                torch._refs._maybe_broadcast = f__maybe_broadcast
-                ShapeEnv._evaluate_expr = f_shape_env__evaluate_expr
-
-                if verbose:
-                    print("[torch_export_patches] restored pytorch functions")
-
-            if patch_torch and stop_if_static:
-                if verbose:
-                    print("[torch_export_patches] restored ShapeEnv._set_replacement")
-
-                ShapeEnv._set_replacement = f_shape_env__set_replacement
-
-                if verbose:
-                    print("[torch_export_patches] restored ShapeEnv._log_guard")
 
-                ShapeEnv._log_guard = f_shape_env__log_guard
-
-                if stop_if_static > 1:
-                    if verbose:
-                        print("[torch_export_patches] restored ShapeEnv._check_frozen")
-                    ShapeEnv._check_frozen = f_shape_env__check_frozen
-
-            if patch_torch and catch_constraints:
-                # to catch or skip dynamic_shapes issues
-                torch._export.non_strict_utils.produce_guards_and_solve_constraints = (
-                    f_produce_guards_and_solve_constraints
-                )
-                torch._export.utils._check_input_constraints_for_graph = (
-                    f__check_input_constraints_for_graph
+            if patch_transformers:
+                _unpatch_transformers(
+                    verbose,
+                    patch_details,
+                    f_transformers__vmap_for_bhqkv,
+                    f_transformers_eager_mask,
+                    f_transformers_sdpa_attention_forward,
+                    f_transformers_sdpa_mask,
+                    f_transformers_sdpa_mask_recent_torch,
+                    revert_patches_info,
                 )
-                if verbose:
-                    print("[torch_export_patches] restored shape constraints")
 
             if custom_patches:
                 if verbose:
@@ -680,118 +1072,6 @@ def torch_export_patches(
                     custom_patches, revert_custom_patches_info, verbose=verbose
                 )
 
-            ##############
-            # transformers
-            ##############
-
-            if patch_transformers:
-                try:
-                    import transformers.masking_utils as masking_utils
-                except ImportError:
-                    masking_utils = None
-                if verbose:
-                    print("[torch_export_patches] unpatches transformers")
-                unpatch_module_or_classes(
-                    patch_transformers_list, revert_patches_info, verbose=verbose
-                )
-
-                if (  # vmap
-                    masking_utils
-                    and patch_transformers_list.patch_masking_utils
-                    and hasattr(masking_utils, "_vmap_for_bhqkv")
-                ):
-                    masking_utils._vmap_for_bhqkv = f_transformers__vmap_for_bhqkv
-
-                    if verbose:
-                        print(
-                            "[torch_export_patches] restored "
-                            "transformers.masking_utils._vmap_for_bhqkv"
-                        )
-
-                    masking_utils.sdpa_mask_recent_torch = (
-                        f_transformers_sdpa_mask_recent_torch
-                    )
-
-                    if verbose:
-                        print(
-                            "[torch_export_patches] restored "
-                            "transformers.masking_utils.sdpa_mask_recent_torch"
-                        )
-
-                    if f_transformers_sdpa_mask is not None:
-                        masking_utils.sdpa_mask = f_transformers_sdpa_mask
-                        if verbose:
-                            print(
-                                "[torch_export_patches] restored "
-                                "transformers.masking_utils.sdpa_mask"
-                            )
-
-                if (  # eager_mask
-                    masking_utils
-                    and patch_transformers_list.patch_masking_utils
-                    and hasattr(masking_utils, "eager_mask")
-                ):
-                    f_transformers_eager_mask = masking_utils.eager_mask
-                    masking_utils.eager_mask = f_transformers_eager_mask
-                    if verbose:
-                        print(
-                            "[torch_export_patches] restored "
-                            "transformers.masking_utils.eager_mask"
-                        )
-                    if (
-                        "eager" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS
-                        and masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["eager"]
-                        == patch_transformers_list.patched_eager_mask
-                    ):
-                        masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["eager"] = (
-                            f_transformers_eager_mask
-                        )
-                        if verbose:
-                            print(
-                                "[torch_export_patches] restored "
-                                "transformers.masking_utils.eager_mask "
-                                "in ALL_MASK_ATTENTION_FUNCTIONS"
-                            )
-
-                if (  # sdpa_mask
-                    masking_utils
-                    and patch_transformers_list.patch_masking_utils
-                    and hasattr(masking_utils, "sdpa_mask")
-                ):
-                    if (
-                        "sdpa" in masking_utils.ALL_MASK_ATTENTION_FUNCTIONS
-                        and masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"]
-                        == patch_transformers_list.patched_sdpa_mask_recent_torch
-                    ):
-                        masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"] = (
-                            f_transformers_sdpa_mask
-                        )
-                        if verbose:
-                            print(
-                                "[torch_export_patches] restored "
-                                "transformers.masking_utils.sdpa_mask "
-                                "in ALL_MASK_ATTENTION_FUNCTIONS"
-                            )
-
-                if (  # sdpa_attention_forward
-                    sdpa_attention is not None
-                    and modeling_utils is not None
-                    and hasattr(sdpa_attention, "sdpa_attention_forward")
-                    and hasattr(sdpa_attention, "use_gqa_in_sdpa")
-                    and hasattr(modeling_utils, "AttentionInterface")
-                ):
-                    sdpa_attention.sdpa_attention_forward = f_sdpa_attention_forward
-                    modeling_utils.sdpa_attention_forward = f_sdpa_attention_forward
-                    modeling_utils.AttentionInterface._global_mapping["sdpa"] = (
-                        f_sdpa_attention_forward
-                    )
-                    if verbose:
-                        print(
-                            "[torch_export_patches] restored "
-                            "transformers.integrations.sdpa_attention."
-                            "sdpa_attention_forward"
-                        )
-
             ########
             # caches
             ########