Implementing the great "Formal Algorithms for Transformers" by Mary Phuong and Marcus Hutter. Purely for my own understanding, don't use this code for Pete's sake!
For ease of implementation in Pytorch, embedded vectors / tensors are transposed w/r/t/ "Formal Algorithms"