torch.matmul does broadcasting for batch dimensions (it does batched matrix-matrix multiply), so no, you don’t need to manually unsqueeze the first two dimensions of the weight tensor. However, we do have to be careful to make sure that the n_head dimension lines up correctly. So if y has shape (batch, windows, n_head, hidden) and weight has shape (n_head, output, hidden), then we would want to line up the n_head dimensions by:

unsqueezing y to have shape (batch, windows, n_head, hidden, 1)

doing matmul(weight, unsqueezed_y) to get an output of shape (batch, windows, n_head, output, 1)

squeezing away the last dimension of size 1.