Patch LlamaRotaryEmbedding

Author: xadupre

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -814,10 +814,7 @@ def wrapper(self, x, position_ids):
     return wrapper
 
 
-class patched_Phi3RotaryEmbedding(torch.nn.Module):
-    _PATCHES_ = ["forward"]
-    _PATCHED_CLASS_ = transformers.models.phi3.modeling_phi3.Phi3RotaryEmbedding
-
+class common_RotaryEmbedding(torch.nn.Module):
     @torch.no_grad()
     @patched_dynamic_rope_update
     def forward(self, x, position_ids):
@@ -843,6 +840,16 @@ def forward(self, x, position_ids):
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
+class patched_LlamaRotaryEmbedding(common_RotaryEmbedding):
+    _PATCHES_ = ["forward"]
+    _PATCHED_CLASS_ = transformers.models.llama.modeling_llama.LlamaRotaryEmbedding
+
+
+class patched_Phi3RotaryEmbedding(common_RotaryEmbedding):
+    _PATCHES_ = ["forward"]
+    _PATCHED_CLASS_ = transformers.models.phi3.modeling_phi3.Phi3RotaryEmbedding
+
+
 class patched_IdeficsEmbedding(torch.nn.Module):
     _PATCHES_ = ["forward"]
     _PATCHED_CLASS_ = transformers.models.idefics.modeling_idefics.IdeficsEmbedding