Upload HyenaDNAForCausalLM

modeling_hyena.py

@@ -60,10 +60,8 @@ class HyenaPositionalEmbedding(nn.Module):
         w = 2 * math.pi * t_rescaled / self.seq_len # 1, L, 1
 
         f = torch.linspace(1e-4, bands - 1, bands)[None, None]
-        # Matt: This is just Euler's formula, so if complex64 is a problem it can be replaced
-        # by separate sin() and cos() calls.
-        z = torch.exp(-1j * f * w)
-        z = torch.cat([t, z.real, z.imag], dim=-1)
+
+        z = torch.cat([t, torch.cos(-f * w), torch.sin(-f * w)], dim=-1)
         # TODO Set z's LR to lr_pos_emb
         self.z = nn.Parameter(z, requires_grad=True)
         self.register_buffer("t", t)