"""可检查形状的因果多头注意力。"""
import math
import torch
from torch import nn


class CausalMHA(nn.Module):
    def __init__(self, dim: int, heads: int):
        super().__init__()
        assert dim % heads == 0
        self.heads, self.head_dim = heads, dim // heads
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.out = nn.Linear(dim, dim, bias=False)

    def forward(self, x):
        b, t, d = x.shape
        q, k, v = self.qkv(x).view(b, t, 3, self.heads, self.head_dim).unbind(2)
        q, k, v = [z.transpose(1, 2) for z in (q, k, v)]
        score = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
        mask = torch.triu(torch.ones(t, t, dtype=torch.bool, device=x.device), 1)
        score = score.masked_fill(mask, torch.finfo(score.dtype).min)
        prob = torch.softmax(score.float(), -1).to(x.dtype)
        context = (prob @ v).transpose(1, 2).contiguous().view(b, t, d)
        return self.out(context), prob


if __name__ == "__main__":
    y, p = CausalMHA(64, 4)(torch.randn(2, 7, 64))
    assert y.shape == (2, 7, 64) and p.shape == (2, 4, 7, 7)
    assert torch.allclose(p[..., 0, 1:], torch.zeros_like(p[..., 0, 1:]))
    print("attention tests passed")