d_k = K.shape[-1]
attention_scores = (Q @ K.T) / math.sqrt(d_k)