Two known answers that test attention

The notebook’s softmax subtracts each row’s maximum first:

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

The test then checks two answers known without running attention:

1. Zero logits average the values uniformly. With 3 keys, every weight is 1/3 and each output is the mean of the values.
2. Logits near 1e4 stay finite. Shifting by the maximum is what keeps exp from overflowing.

Scores are also checked against an explicit-loop oracle — one np.dot per batch, head, query and key.