Possible mistakes in d_k, d_v of MultiheadAttention

In the __init__ function of the MultiheadAttention class you use d_k and d_v to denote the dimensions of keys and values. You also define the projections below:

self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)

However, when d_v is not the same as d_q (it should be d_model // n_head), this will cause the shape of the queries to change after the attention operation and will cause problems in multiple layer structures. After going through the official MultiheadAttention implementation of PyTorch, I believe that you used a similar presentation with:

self.q_proj_weight = Parameter(torch.empty((embed_dim, embed_dim), **factory_kwargs))
self.k_proj_weight = Parameter(torch.empty((embed_dim, self.kdim), **factory_kwargs))
self.v_proj_weight = Parameter(torch.empty((embed_dim, self.vdim), **factory_kwargs))

However, in the official PyTorch implementation, it used weights, rather than a nn.Linear class, which means that the weights are actually used to transformer the dimension of keys from self.kdim to embed_dim, which is the very opposite to what your implementation is doing. So I believe that there might be some errors with your code. But overall, thank you for your work, it helped me a lot.