A transformer is the same block stacked many times. Each block has two steps: attention mixes information between positions, then a small feed-forward network works on each position alone.

Two details make deep stacks trainable:

• Residual connections: each step adds its result to its input, x = x + step(x), so information and gradients pass straight through.
• Layer normalisation before each step keeps the numbers in a steady range.

import torch
from torch import nn

ln = nn.LayerNorm(8)
x = torch.randn(4, 8) * 50 + 3
y = ln(x)
print(y.mean(-1), y.std(-1))
# about 0 and 1 for each row

The feed-forward part widens to a hidden size, applies GELU, and narrows back.

Write Block(d, hidden): an nn.Module with ln1 and ln2 (nn.LayerNorm(d)), attn (the given SelfAttention(d)) and mlp, an nn.Sequential of Linear(d, hidden), GELU and Linear(hidden, d). forward(x) adds attention of ln1(x) to x, then adds the MLP of ln2 of the result.

The explanation opens here once every check passes.