实现Masked Multi-Head Self-Atten_牛客题霸

给定批量序列表示 X（形状：[batch, seq, d_model]）与权重矩阵 W_Q、W_K、W_V、W_O（均为 d_model×d_model），实现 Masked Multi-Head Self-Attention。将最后一维按头数 num_heads 均分，每头维度 d_k = d_model num_heads。计算步骤： 1) Q = X @ W_Q，K = X @ W_K，V = X @ W_V。 2) 将 Q、K、V reshape 为 [batch, num_heads, seq, d_k]。 3) 计算注意力分数 scores = (Q @ K^T) sqrt(d_k)，其中 K^T 表示每头在最后两维做转置得到 [batch, num_heads, seq, seq]。 4) 使用下三角因果掩码（只能看见当前及更早位置）：掩掉上三角元素（置为一个很小的负数）。 5) 在最后一维做 softmax 得到权重，注意数值稳定性（减去每行最大值再做 exp）。 6) attention = softmax @ V（形状 [batch, num_heads, seq, d_k]）。 7) 拼回 [batch, seq, d_model] 后，再右乘 W_O。输出保留两位小数，结果需转换为 Python List。

输入描述:

以分号分隔的 6 个参数：num_heads; X; W_Q; W_K; W_V; W_O  其中 X、W_Q、W_K、W_V、W_O 用 Python 风格的嵌套列表表示。

输出描述:

最终输出张量（形状 [batch, seq, d_model]），四舍五入到小数点后两位，类型为 List。

示例1

输入

2; [[[1, 1], [1, 1], [1, 1]]]; [[1, 0], [0, 1]]; [[1, 0], [0, 1]]; [[1, 0], [0, 1]]; [[1, 0], [0, 1]]

输出

[[[1.00, 1.00], [1.00, 1.00], [1.00, 1.00]]]

说明

权重为单位矩阵，Q=K=V=X。因果掩码使第 i 个位置只看见前 i+1 个位置；由于各位置完全相同，softmax 权重在可见范围内均匀分布，输出与输入一致；乘 W_O（单位）后不变。

备注:

本题由牛友@Charles 整理上传