[ET-VK][ez] Apply rotary embedding as Module

examples/models/llama/llama_transformer.py

@@ -15,10 +15,10 @@
 import torch.nn.functional as F
 
 from executorch.examples.models.llama.rope import (
-    apply_rotary_emb,
     hf_apply_rotary_emb,
     hf_precompute_freqs_cis,
     precompute_freqs_cis,
+    RotaryEmbedding,
 )
 
 from torch import nn
@@ -311,7 +311,7 @@ def __init__(self, args: ModelArgs, layer_id: int):
         if args.use_hf_rope:
             self.apply_rotary_emb = hf_apply_rotary_emb
         else:
-            self.apply_rotary_emb = apply_rotary_emb
+            self.apply_rotary_emb = RotaryEmbedding()
 
     def forward(
         self,

examples/models/llama/rope.py

@@ -92,6 +92,21 @@ def apply_rotary_emb(
     return xq_out.type_as(xq), xk_out.type_as(xk)
 
 
+class RotaryEmbedding(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(
+        self,
+        xq: torch.Tensor,
+        xk: torch.Tensor,
+        freqs_cos: torch.Tensor,
+        freqs_sin: torch.Tensor,
+    ):
+        xq_out, xk_out = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)
+        return xq_out, xk_out
+
+
 # ======================= HuggingFace Implementation ========================