Why aren't you transposing the input in multi-head attention? about machine-learning-collection HOT 1 CLOSED

    values = self.values(values)  # (N, value_len, embed_size)
    keys = self.keys(keys)  # (N, key_len, embed_size)
    queries = self.queries(query)  # (N, query_len, embed_size)

    # Split the embedding into self.heads different pieces
    values = values.reshape(N, value_len, self.heads, self.head_dim)
    keys = keys.reshape(N, key_len, self.heads, self.head_dim)
    queries = queries.reshape(N, query_len, self.heads, self.head_dim)

    my_keys, my_values, my_queries = keys.transpose(1,2), values.transpose(1,2), queries.transpose(1,2)
    my_energy = torch.matmul(my_queries, my_keys.transpose(2,3))

    # Einsum does matrix mult. for query*keys for each training example
    # with every other training example, don't be confused by einsum
    # it's just how I like doing matrix multiplication & bmm

    energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
    assert torch.equal(my_energy, energy), "Implementations are different"