attention_scores = queries @ keys.T
attention_weights = torch.softmax(
    attention_scores / keys.shape[-1] ** 0.5,
    dim=-1,
)
context_vectors = attention_weights @ values

        __
        __