Masked Attention

Task

This is a small transformer written with plain Python lists: an embedding, one causal self-attention layer with two heads, and a projection to scores over a six-token vocabulary. Its weights are fixed. Some requests return NaN, and the tokens it samples are not the ones the same weights give in the reference implementation. The weights are not the problem. Find the causes and fix them. Do not hide them.

Requirements

1. transformer(request) takes one request object and returns a small object. tokens is a list of 1 to 8 integers from 0 to 5 (ValueError containing tokens otherwise). Any other op raises ValueError containing op. Do not change the arguments.
2. {op: "mask", length} returns {mask}: a length by length table where mask[query][key] is 1 when the query position may look at the key position and 0 when it may not. A position may look at itself and at earlier positions, never at later ones. length is an integer from 1 to 8 (ValueError containing length otherwise).
3. {op: "weights", tokens} returns {weights}: weights[head][query][key], the attention weights of each of the two heads. Every query spreads a total weight of 1 over the keys it may look at, and gives 0 to the rest.
4. {op: "logits", tokens} returns {logits}: one row of six scores per position. The scores at a position depend only on the tokens at that position and before it.
5. {op: "generate", tokens, steps} returns {tokens}: the input followed by steps greedy picks, each the highest-scoring next token (ties go to the lower token). steps is a non-negative integer (ValueError containing steps otherwise), and tokens plus steps must fit in 8 (ValueError containing length otherwise).
6. Every number in a result is finite and rounded to 6 decimals with the rounded() helper. The weight tables at the top of the file must stay as they are.