X = ["the", "cat", "sat"]
Attention(X) == Attention(["sat", "cat", "the"])

import torch
import matplotlib.pyplot as plt
import numpy as np

def plot_sinusoidal_pe(d_model=16, max_pos=20):
    pos = torch.arange(max_pos).unsqueeze(1)
    i = torch.arange(0, d_model, 2)
    div_term = torch.exp(i * -torch.log(torch.tensor(10000.0)) / d_model)
    pe_even = torch.sin(pos * div_term)
    pe_odd = torch.cos(pos * div_term)

    pe = torch.zeros(max_pos, d_model)
    pe[:, 0::2] = pe_even
    pe[:, 1::2] = pe_odd

    plt.figure(figsize=(12, 6))
    for dim in range(0, d_model, 2):
        plt.plot(pos, pe[:, dim], label=f"dim {dim}" if dim < 6 else "")
    plt.legend()
    plt.xlabel("Position")
    plt.ylabel("PE Value")
    plt.title("Sinusoidal PE: Different Frequencies per Dimension")
    plt.grid(True, alpha=0.3)
    plt.show()

plot_sinusoidal_pe()

# Relative distance encoding
pos_i, pos_j = 5, 10
pe_i = pe[pos_i]
pe_j = pe[pos_j]

# Dot product peaks at fixed relative distance
dist = 5
correlations = []
for offset in range(-10, 11):
    if 0 <= pos_i + offset < max_pos:
        corr = torch.dot(pe[pos_i], pe[pos_i + offset])
        correlations.append((offset, corr.item()))

offsets, corrs = zip(*correlations)
plt.plot(offsets, corrs, 'o-')
plt.title("PE Correlation vs Relative Position")
plt.xlabel("Position Offset")
plt.ylabel("Dot Product")
plt.show()

from sklearn.metrics.pairwise import cosine_similarity

pos1, pos2 = 100, 105
pe1 = pe[pos1].unsqueeze(0)
pe2 = pe[pos2].unsqueeze(0)
sim = cosine_similarity(pe1.numpy(), pe2.numpy())[0][0]
print(f"Cosine sim(pos=100, 105) = {sim:.3f}")  # ~0.999

class LearnedPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):
        super().__init__()
        self.pe = nn.Embedding(max_len, d_model)
        nn.init.normal_(self.pe.weight, std=0.02)

    def forward(self, x):
        seq_len = x.size(1)
        pos = torch.arange(seq_len, device=x.device)
        return x + self.pe(pos)

import torch.optim as optim

def train_copy_task(model_cls, use_learned_pe=False, max_len=20):
    model = nn.Sequential(
        nn.Embedding(10, 16),
        model_cls(d_model=16, num_heads=4, use_learned_pe=use_learned_pe),
        nn.Linear(16, 10)
    )
    opt = optim.Adam(model.parameters(), lr=0.01)
    criterion = nn.CrossEntropyLoss()

    losses = []
    for epoch in range(300):
        src = torch.randint(0, 5, (32, max_len))
        tgt = src.clone()

        logits = model[0](src)
        logits = model[1](logits)[0]
        logits = model[2](logits)

        loss = criterion(logits.view(-1, 10), tgt.view(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
        losses.append(loss.item())

    return losses

# Run both
loss_sine = train_copy_task(TransformerBlock, use_learned_pe=False)
loss_learned = train_copy_task(TransformerBlock, use_learned_pe=True)

plt.plot(loss_sine, label="Sinusoidal PE")
plt.plot(loss_learned, label="Learned PE")
plt.legend()
plt.title("Copy Task: Sinusoidal vs Learned PE")
plt.xlabel("Training Step")
plt.ylabel("Loss")
plt.show()

# Train on max_len=20
model_sine = ...  # trained with sinusoidal
model_learned = ...  # trained with learned (max_len=20)

# Test on length 50
long_seq = torch.randint(0, 5, (1, 50))
with torch.no_grad():
    out_sine = model_sine(long_seq)
    # out_learned → IndexError! (Embedding size = 20)

def apply_rotary_emb(q, k, freqs):
    # q, k: (B, H, N, d_k)
    q_real, q_imag = q[..., :d_k//2], q[..., d_k//2:]
    k_real, k_imag = k[..., :d_k//2], k[..., d_k//2:]

    # Rotate
    q_rot = torch.cat([-q_imag, q_real], dim=-1) * freqs
    k_rot = torch.cat([-k_imag, k_real], dim=-1) * freqs

    return q_rot + q, k_rot + k

pe_sine = SinusoidalPositionalEncoding(128, 100).pe[0].cpu().numpy()
pe_learned = LearnedPositionalEncoding(128, 100).pe.weight.detach().cpu().numpy()

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
sns.heatmap(pe_sine, ax=ax1, cmap="RdYlBu", center=0)
sns.heatmap(pe_learned, ax=ax2, cmap="RdYlBu", center=0)
ax1.set_title("Sinusoidal PE")
ax2.set_title("Learned PE (Random Init)")
plt.show()

import torch
import torch.nn as nn

# === Sinusoidal ===
class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * -(torch.log(torch.tensor(10000.0)) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer('pe', pe.unsqueeze(0))
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

# === Learned ===
class LearnedPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):
        super().__init__()
        self.pe = nn.Embedding(max_len, d_model)
    def forward(self, x):
        pos = torch.arange(x.size(1), device=x.device)
        return x + self.pe(pos)