restrict or optimize comparisons

for i in tokens:
    for j in selected_tokens:
        score[i][j] = dot(Q[i], K[j])