import numpy as np

def softmax(x):
    x = x - x.max(axis=-1, keepdims=True)          # 数值稳定：先减最大值
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V, causal=False):
    d_k = Q.shape[-1]
    scores = Q @ K.swapaxes(-2, -1) / np.sqrt(d_k)   # (T, T)
    if causal:
        T = scores.shape[-1]
        mask = np.triu(np.ones((T, T), dtype=bool), k=1)  # 对角线以上 = 未来
        scores = np.where(mask, -np.inf, scores)
    weights = softmax(scores)                          # 每行和为 1
    return weights @ V, weights                        # (T, d_v), (T, T)

rng = np.random.default_rng(0)
T, d_model, d_k, d_v = 4, 8, 16, 16
X = rng.standard_normal((T, d_model))
W_q, W_k, W_v = (rng.standard_normal((d_model, d)) / np.sqrt(d_model) for d in (d_k, d_k, d_v))
Q, K, V = X @ W_q, X @ W_k, X @ W_v

out, w = attention(Q, K, V, causal=True)
print("Q", Q.shape, "K", K.shape, "V", V.shape)
print("out", out.shape, "weights", w.shape)
print("row sums:", w.sum(axis=-1).round(3))
print("causal weights:\n", w.round(2))

# 实验：为什么要除以 sqrt(d_k)
for d in (16, 512):
    q, k = rng.standard_normal((2, 10000, d))
    dots = (q * k).sum(-1)
    print(f"d_k={d:4d}  var(q·k)={dots.var():7.1f}  var(q·k/sqrt(d_k))={(dots/np.sqrt(d)).var():.2f}")
S = rng.standard_normal((1000, 8)) * np.sqrt(512)     # 1000 行未缩放分数，每行 8 个 key
print("avg max weight  unscaled:", softmax(S).max(-1).mean().round(3),
      " scaled:", softmax(S / np.sqrt(512)).max(-1).mean().round(3))