def forward(self, queries, keys, values, valid_lens):
    queries = self.transpose_qkv(np.dot(queries, self.W_q) + self.b_q)
    keys = self.transpose_qkv(np.dot(keys, self.W_k) + self.b_k)
    values = self.transpose_qkv(np.dot(values, self.W_v) + self.b_v)

    if valid_lens is not None:
        valid_lens = np.repeat(valid_lens, self.num_heads, axis=0)

    output = self.scaled_dot_product_attention(queries, keys, values, valid_lens)
    output_concat = self.transpose_output(output)
    return np.dot(output_concat, self.W_o) + self.b_o