A Transformer Block in Python

A transformer block is the unit that real models stack dozens of times. It has two parts. First attention: each position gathers what it needs from the positions before it. Then a feed-forward network, like your digits network: 32 numbers to 128, ReLU, back to 32, run on each position by itself to work on what attention gathered.

Before each part sits an nn.LayerNorm(32). It rescales the 32 numbers of each position so they average 0 with a spread of 1 and gives back a tensor of the same shape, so the part after it always gets numbers of a steady size:

norm = nn.LayerNorm(3)
v = torch.tensor([2.0, 4.0, 6.0])
print(norm(v).tolist())   # about [-1.22, 0.0, 1.22]

After each part comes a residual connection: the part's answer is added to x instead of replacing it. If a part has nothing useful to say yet, x passes through unchanged, which keeps a tall stack of blocks easy to train. Two more tools: nn.Sequential(a, b, c) runs layers one after the other and gives back the last one's output, and nn.ReLU() is ReLU as a layer.

Run on a batch, your code and the checks after it print:

h: [32, 8, 32] out: [32, 8, 32]
after norm1, one position: mean 0.0 spread 1.0
attention and feed-forward silenced, the block gives back h: True
the block changes h by about: 0.288
change position 7: positions 0-6 stay the same: True
weights in one block: 12608
three blocks in a row: [32, 8, 32]

Your task: the parts from the last lessons are ready, and Block's forward so far just gives x back. In __init__, store an nn.LayerNorm(n_embd) in self.norm1, a MultiHead() in self.attention, a second LayerNorm in self.norm2 and the feed-forward nn.Sequential in self.feed_forward. In forward, above return x, add attention's answer to x, then the feed-forward's, each part reading its LayerNorm of x.

Assignment
Follow the Coding Tutorial to build the four parts of the block and its two residual lines.

Hint
Both lines of forward have the same shape: x becomes x plus the part, run on the norm of x. Write the 128 inside the feed-forward as 4 * n_embd.