Linear Layer Forward & Backward

A linear (fully connected) layer is the most common block in an LLM: the Q/K/V projections, the MLP and the final vocabulary head are all linear layers. Here you write its forward pass and derive its backward pass by hand.

For a batch of $N$ inputs:

Y = XW + b, \qquad X \in \mathbb{R}^{N \times D_{in}},\; W \in \mathbb{R}^{D_{in} \times D_{out}},\; b \in \mathbb{R}^{D_{out}}

Implement:

linear_forward(x, W, b) returns $Y$ .
linear_backward(x, W, b, grad_out): given the upstream gradient $\frac{\partial L}{\partial Y}$ (grad_out, shape $N \times D_{out}$ ), return the tuple (dx, dW, db) with the gradients of $L$ with respect to $X$ , $W$ and $b$ .

Requirements

Each gradient has the same shape as its parameter. Use this to sanity-check your formulas.
No loops: three matrix expressions are enough.

Hints

Work out one element first: $Y_{nk} = \sum_d X_{nd} W_{dk} + b_k$ . Which $Y$ entries does $W_{dk}$ touch?
$b$ is broadcast over the batch. In the backward pass, broadcasting becomes a sum.

The tests check your gradients against finite differences, $\frac{L(\theta + \epsilon) - L(\theta - \epsilon)}{2\epsilon}$ . That is how you debug any hand-written backward pass.