patch

Author: xadupre

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -350,13 +350,14 @@ def torch_export_patches(
                 patch_transformers_list, verbose=verbose
             )
 
-        if 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 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 custom_patches:
             if verbose:
@@ -450,6 +451,10 @@ def torch_export_patches(
             ##############
 
             if patch_transformers:
+                try:
+                    import transformers.masking_utils as masking_utils
+                except ImportError:
+                    masking_utils = None
                 if verbose:
                     print("[torch_export_patches] unpatch transformers")
                 unpatch_module_or_classes(

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -13,7 +13,10 @@ def patched__vmap_for_bhqkv(mask_function: Callable, bh_indices: bool = True) ->
     """Patch for function ``transformers.masking_utils._vmap_for_bhqkv``."""
     from ...helpers import string_type
 
-    dimensions: List[Optional[int]] = [(None, None, None, 0), (None, None, 0, None)]
+    dimensions: List[Tuple[Optional[int], ...]] = [
+        (None, None, None, 0),
+        (None, None, 0, None),
+    ]
     if bh_indices:
         dimensions.extend([(None, 0, None, None), (0, None, None, None)])
     dimensions = [tuple(1 if d is None else -1 for d in shape) for shape in dimensions]