AI006. 因果注意力计算

# 因果注意力:第 i 行只对 j<=i 打分 dot(Q[i],K[j])/√d,行内减 max 再 softmax,加权和 V;
# 掩码用收窄循环范围实现,max 也只在掩码内取;输出四位小数注意负零钳制。
import math
import sys


def solve() -> None:
    data = list(map(float, sys.stdin.buffer.read().split()))
    if not data:
        return
    at = 0
    n = int(data[at]); at += 1
    d = int(data[at]); at += 1
    def read_matrix():
        nonlocal at
        matrix = []
        for _ in range(n):
            matrix.append(data[at:at + d])
            at += d
        return matrix
    q = read_matrix()
    k = read_matrix()
    v = read_matrix()
    scale = math.sqrt(d)
    ans = []
    for i in range(n):
        scores = []
        for j in range(i + 1):
            scores.append(sum(q[i][p] * k[j][p] for p in range(d)) / scale)
        mx = max(scores)
        weights = [math.exp(score - mx) for score in scores]
        total = sum(weights)
        row = [0.0] * d
        for j, weight in enumerate(weights):
            weight /= total
            for p in range(d):
                row[p] += weight * v[j][p]
        ans.append(" ".join(f"{0.0 if abs(value) < 0.00005 else value:.4f}" for value in row))
    print("\n".join(ans))


if __name__ == "__main__":
    solve()