P4283.第2题-JSON数组
题目内容
请在仅使用 numpy 的前提下，实现单头 Scaled−Dot Product Attention 的前向计算。
输入为 Query(Q),Key(K)，Value(V) 三矩阵和可选 attention_mask 。
设 dk=d 。
1.缩放点积得分 S=dkQKT
2.加掩码(如果提供) S←S+mask
- 掩码元素为 0 或 −1e9 ——等价于将被屏蔽位置的 softmax 概率压至 0
3.行级数值稳定 softmax
S=S−S.max(axis=1,keepdims=True)
P=exp(S)/exp(S).sum(axis=1,keepdims=True)
4.加权求和 O=PV ，结果形状 (m,dv)
5.输出展开
- 按行优先顺序进行展平操作，即将 O[0],O[1],...,O[m−1] 依次连接成一个一维数组
输入描述
单行 JSON：
{
"Q":[[...],...], //形状 (m,d)
"K":[[...],...], //形状 (n,d)
"V":[[...],...], //形状 (n,dv)；本题 dv==d
"mask":[[...],...] or null //形状 (m,n) ；0 表示可见，−1e9 表示屏蔽
}
-
2≤d≤4,1≤m,n≤3
-
所有值为实数
-
若"mask"为 null ，视为全 0
输出描述
使用 round(x,6) 保留小数位即可，返回单行 JSON 数组
补充说明
为了确保通过测试用例，仅允许使用 Numpy 实现
样例1
输入
{"Q":[[1,0]],"K":[[1,0]],"V":[[5,5]],"mask":null}
输出
[5.0, 5.0]
