From 40169f9fd3cff63cfbee1fcfe9841505d60f3833 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Mon, 26 Feb 2024 04:23:38 +1100
Subject: [PATCH] Update gemma.py

---
 unsloth/models/gemma.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py
index ac7a3136c1..0107fcf8d3 100644
--- a/unsloth/models/gemma.py
+++ b/unsloth/models/gemma.py
@@ -97,6 +97,18 @@ class FastGemmaRotaryEmbedding(torch.nn.Module):
             self.cos_cached = emb.cos().to(dtype=x.dtype)
             self.sin_cached = emb.sin().to(dtype=x.dtype)
         pass
+
+        position_ids = torch.arange(self.max_position_embeddings, device=x.device, dtype=torch.int64).unsqueeze(0)
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(1, -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos_cached = emb.cos().to(dtype=x.dtype)
+        sin_cached = emb.sin().to(dtype=x.dtype)
+
+        print(cos_cached)
+        print(self.cos_cached)
+
         # return self.cos_cached[:,:length], self.sin_cached[:,:length]
         return self.cos_cached, self.sin_cached
 pass