PyTorch MultiheadAttention

See this first:

• https://pytorch.org/tutorials/intermediate/transformer_building_blocks.html

Transformers building blocks.

https://pytorch.org/tutorials/prototype/nestedtensor.html#why-nested-tensor