import torch

import torch.nn.functional as F

def self_attention(Q, K, V, d_k):

   # 步骤1：计算注意力分数

   scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))

   # 步骤2：softmax归一化

   attn_weights = F.softmax(scores, dim=-1)

   # 步骤3：加权求和得到输出

   output = ______________

   return output

# 测试输入

Q = torch.randn(2, 5, 64)

K = torch.randn(2, 5, 64)

V = torch.randn(2, 5, 64)

d_k = 64