https://d2l.ai/chapter_attention-mechanisms-and-transformers/self-attention-and-positional-encoding.html