Dos respuestas conocidas que ponen a prueba la atención

El softmax del cuaderno resta primero el máximo de cada fila:

shifted = logits - logits.max(axis=-1, keepdims=True)
weights = np.exp(shifted); weights /= weights.sum(axis=-1, keepdims=True)

La prueba comprueba luego dos respuestas conocidas sin ejecutar la atención:

1. Los logits cero promedian los valores de forma uniforme. Con 3 claves, cada peso es 1/3 y cada salida es la media de los valores.
2. Los logits cercanos a 1e4 siguen siendo finitos. Restar el máximo es lo que evita que exp se desborde.

Las puntuaciones también se comparan con un oráculo de bucles explícitos: un np.dot por lote, cabeza, consulta y clave.