Hi!
I’m trying to understand the multiheadattention function at pytorch MultiheadAttention — PyTorch 1.8.1 documentation 3 and if I can use it to compute the selfattention as it can be done in the keras implementation MultiHeadAttention layer 5
With Keras implementation I’m able to run selfattention over a 1D vector the following way:
import tensorflow as tf

layer = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=2)
input_tensor = tf.keras.Input(shape=[8, 16])
output_tensor = layer(input_tensor, input_tensor)
print(output_tensor.shape)
(None, 8, 16)

I’ve tried to do the same with the pytorch implementation but couldn’t make it work. Could be because of lacking knowledge of the Transformer implementation.
import torch
import torch.nn as nn
multihead_attn = nn.MultiheadAttention(16, 8)
input_tensor = torch.zeros(8,60)
multihead_attn(input_tensor,input_tensor,input_tensor)

With error
-> 4624     tgt_len, bsz, embed_dim = query.size()
   4625     assert embed_dim == embed_dim_to_check
   4626     # allow MHA to have different sizes for the feature dimension

ValueError: not enough values to unpack (expected 3, got 2)

It seems the keras function does some work under the hood to make it easier for the user and perform self-attention when both Query and Value are similar. Is there a way to do the same with pytorch?