# 縮放點積注意力：照公式自己寫，和 PyTorch 內建的版本比對；並看為什麼要除以 √d
import math
import torch
import torch.nn.functional as F

torch.manual_seed(0)
n, d = 5, 16                                       # 5 個位置，每個向量 16 維
Q, K, V = torch.randn(n, d), torch.randn(n, d), torch.randn(n, d)

def attention(Q, K, V, causal=False):
    S = Q @ K.T / math.sqrt(K.shape[-1])           # 每個 query 與每個 key 的相似度
    if causal:                                     # 因果遮罩：只能看自己與之前的位置
        S = S.masked_fill(torch.ones(n, n).triu(1).bool(), float("-inf"))
    A = torch.softmax(S, dim=-1)                   # 每一列加總為 1：注意力權重
    return A @ V, A

out, A = attention(Q, K, V)
print("自己寫的 = PyTorch 內建：", torch.allclose(out, F.scaled_dot_product_attention(Q, K, V), atol=1e-6))
out_c, A_c = attention(Q, K, V, causal=True)
print("因果遮罩也相同：", torch.allclose(out_c, F.scaled_dot_product_attention(Q, K, V, is_causal=True), atol=1e-6))
print("因果遮罩的權重（上三角為 0）：\n", A_c.round(decimals=2))

# 為什麼除以 √d：分量為標準常態時，q·k 的變異數等於 d，d 越大 softmax 越接近 one-hot
for d in (4, 64, 512):
    q, k = torch.randn(10000, d), torch.randn(10000, d)
    dots = (q * k).sum(1)
    S = torch.randn(1000, 8, d) @ torch.randn(1000, d, 8)          # 1000 組、每組 8×8 的分數
    top = lambda s: torch.softmax(s, -1).max(-1).values.mean().item()
    print(f"d = {d:3d}：q·k 變異數 {dots.var().item():6.1f}，最大權重平均 不縮放 {top(S):.2f}／除以 √d {top(S / math.sqrt(d)):.2f}")