Look Back, Never Ahead in Python

A language model guesses the next character at every position of a window at once. That only works if a position can't see the ones after it: otherwise it would simply read the answer. In this lesson you build the trick that stops it, on a tiny example: 4 positions with 2 numbers each, where every position takes the average of itself and the positions before it.

The rules go in one table. torch.ones(T, T) gives back a T by T table of ones, and torch.tril(t) gives back its lower triangle: the diagonal and what is below it stay, everything above becomes 0. Row t then has ones exactly at the positions t may see.

Then three moves. scores.masked_fill(mask, value) gives back a copy of scores with value wherever mask is True: put float("-inf"), minus infinity, wherever the triangle is 0. torch.softmax(scores, dim=1) turns each row into weights that add up to 1: equal scores get equal weights, and minus infinity gets exactly 0. Finally weights @ x gives each position the weighted sum of the rows of x.

m = torch.tril(torch.ones(3, 3))
print(m.tolist())   # [[1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]
s = torch.zeros(2).masked_fill(torch.tensor([False, True]), float("-inf"))
print(torch.softmax(s, dim=0).tolist())   # [1.0, 0.0]

Here every score starts at 0, so the weights come out equal. In the next lesson the model works out the scores itself, and the mask stays exactly the same. When it's done, your script prints the triangle, then each position's scores, weights and output, then two checks:

tril:
[1.0, 0.0, 0.0, 0.0]
[1.0, 1.0, 0.0, 0.0]
[1.0, 1.0, 1.0, 0.0]
[1.0, 1.0, 1.0, 1.0]
position 0 | scores [0.0, -inf, -inf, -inf] | weights [1.0, 0.0, 0.0, 0.0] | out [2.0, 1.0]
position 1 | scores [0.0, 0.0, -inf, -inf] | weights [0.5, 0.5, 0.0, 0.0] | out [3.0, 2.0]
position 2 | scores [0.0, 0.0, 0.0, -inf] | weights [0.333, 0.333, 0.333, 0.0] | out [4.0, 3.0]
position 3 | scores [0.0, 0.0, 0.0, 0.0] | weights [0.25, 0.25, 0.25, 0.25] | out [5.0, 4.0]
change the last position: the first three outputs stay the same: True
without the mask, position 0 would get [5.0, 4.0] from positions it must not see

Your task: store the lower triangle of a 4 by 4 table of ones in tril, put minus infinity into scores wherever tril is 0, softmax each row of scores into weights, and multiply weights by x into out.

Assignment
Follow the Coding Tutorial to make every position average only itself and the positions before it.

Hint
tril == 0 is True above the diagonal, exactly where the scores must become float("-inf"). Softmax works row by row with dim=1.