[compile] workaround for compile error FakeTensor no op for builtin.- (#2891)

Author: IvanKobzarev

torchtune/models/llama4/_position_embeddings.py

@@ -179,10 +179,14 @@ def forward(
         # tensor has shape [b, s, n_h, h_d // 2, 2]
         x_out = torch.stack(
             [
-                xshaped[..., 0] * rope_cache[..., 0]
-                - xshaped[..., 1] * rope_cache[..., 1],
-                xshaped[..., 1] * rope_cache[..., 0]
-                + xshaped[..., 0] * rope_cache[..., 1],
+                torch.sub(
+                    xshaped[..., 0] * rope_cache[..., 0],
+                    xshaped[..., 1] * rope_cache[..., 1],
+                ),
+                torch.add(
+                    xshaped[..., 1] * rope_cache[..., 0],
+                    xshaped[..., 0] * rope_cache[..., 1],
+                ),
             ],
             -1,
         )