nn.MultiheadAttention的输入输出

nn.MultiheadAttention是PyTorch中的一个模块，用于实现多头注意力机制。它的输入和输出如下所示：

输入：
- query: shape为(batch_size, seq_len, embed_dim)，表示查询序列的张量。
- key: shape为(batch_size, seq_len, embed_dim)，表示键序列的张量。
- value: shape为(batch_size, seq_len, embed_dim)，表示值序列的张量。
- key_padding_mask: shape为(batch_size, 1, seq_len)，表示键序列的填充掩码，用于屏蔽填充位置的注意力权重。
- attn_mask: shape为(seq_len, seq_len)，表示注意力掩码，用于屏蔽无效的注意力权重。

输出：
- output: shape为(batch_size, seq_len, embed_dim)，表示多头注意力机制的输出张量。
- attn_output_weights: shape为(batch_size, num_heads, seq_len, seq_len)，表示注意力权重张量。