Error: la forma correcta implica cabezas correctas

«Puedo aplicar reshape directamente a (B, H, S, D_k).»

example = np.arange(24).reshape(2, 3, 4)
correct = example.reshape(2, 3, 2, 2).transpose(0, 2, 1, 3)
wrong   = example.reshape(2, 2, 3, 2)
assert wrong.shape == correct.shape
assert not np.array_equal(wrong, correct)

Ambos arreglos tienen la misma forma. Solo reshape seguido de transpose mantiene las características de un token dentro de una cabeza. Un reshape directo agrupa lo que esté contiguo en memoria, así que la cabeza 0 del token 1 debería ser [4, 5] y el reshape directo le da las características de otro token.

Sigue un valor para exponer la diferencia: una prueba de forma no puede. Los ejes unitarios pueden ocultar el error, porque con un solo token o una sola cabeza los dos órdenes coinciden.