Por qué (B, H) va primero

Las dos últimas dimensiones son los ejes de la matriz. Poner (B, H) primero permite que la multiplicación de matrices procese en lote ambos ejes, así que un solo @ o einsum calcula todas las cabezas de todos los ejemplos a la vez.

Cada cabeza compara luego cada token consulta con cada token clave en su propio espacio de características de D_k características.

El reshape no aprende esos espacios: los aprenden las matrices de proyección. Separar las cabezas solo reagrupa características que W_Q, W_K y W_V ya produjeron.