软注意力机制代码

以下是一个简单的软注意力机制的代码实现，其中涉及到两个输入矩阵A和B，输出矩阵C的计算过程：

import torch

def soft_attention(A, B):
    # 计算注意力分数
    score = torch.matmul(A, B.transpose(1, 0))
    # 计算注意力权重
    attn_weight = torch.softmax(score, dim=1)
    # 计算注意力加权矩阵
    attn_output = torch.matmul(attn_weight, B)
    # 返回注意力加权矩阵
    return attn_output

在上述代码中，`matmul`函数表示矩阵乘法，`transpose`函数表示矩阵转置，`softmax`函数表示对输入进行softmax操作。注意力分数矩阵的维度为(A的行数, B的行数)，注意力权重矩阵的维度为(A的行数, B的行数)，注意力加权矩阵的维度为(A的行数, B的列数)。