Attention 伪代码实现（pytorch 版）

Attention的原理已经有很多介绍了，实现的伪代码参照transformer，下面写了最简单的版本

import torch, math
from torch import nn
dropout_prob = 0.1

def forward(
        hidden_size, # d
        input, #(b, s, d)
        attention_mask  #(b, s, s)
):
    query = nn.Linear