请按照下方说明实现 multi_head_attention()。
所有 head 的 head_dim 相同；seq ≥ 1。
单头恒等映射
正确性调用
multi_head_attention(tensor[1, 1, 2] [1,0], tensor[1, 1, 2] [1,0], tensor[1, 1, 2] [5,7])
预期结果
tensor[1, 1, 2] [5,7]
双 Token 聚焦
正确性调用
multi_head_attention(tensor[1, 2, 2] [1,0,0,1], tensor[1, 2, 2] [1,0,0,1], tensor[1, 2, 2] [10,0,0,20])
预期结果
tensor[1, 2, 2] [6.697615,6.604769,3.302385,13.395231]
运行会检查可见示例，提交会运行完整的隐藏测试。
登录并提交后，你的提交记录会显示在这里。