# 加分题：用 PyTorch 官方实现对拍（pip install torch）
import math, torch
import torch.nn.functional as F

torch.manual_seed(0)
B, H, T, d_k = 2, 4, 5, 16                       # batch, heads, seq_len, head_dim
q, k, v = torch.randn(3, B, H, T, d_k).unbind(0)

def my_attention(q, k, v):
    scores = q @ k.transpose(-2, -1) / math.sqrt(k.size(-1))          # (B, H, T, T)
    causal = torch.tril(torch.ones(T, T, dtype=torch.bool))
    scores = scores.masked_fill(~causal, float("-inf"))
    return F.softmax(scores, dim=-1) @ v                              # (B, H, T, d_k)

mine = my_attention(q, k, v)
ref = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print("shape:", tuple(mine.shape))
print("allclose:", torch.allclose(mine, ref, atol=1e-6))