Linear bwd

The problem, from first principles

[01-linear-forward] computed y = x @ weight.T + bias. To actually train a network, gradient descent needs to know how a small change in each of x, weight, and bias would change the final loss, given only the gradient of the loss with respect to y (the "upstream gradient," handed down from whatever layer comes after this one). [02-deep-learning-core/04-autograd]'s matmul_backward already derived the general rule for backpropagating through a matrix multiplication; this question applies that exact rule to the specific shapes a linear layer uses.

Three different gradients are needed here, for three different reasons: grad_x is needed so the PREVIOUS layer can continue backpropagating further back through the network; grad_weight and grad_bias are needed directly by the optimizer (every optimizer question in this Part, sgd_step, adam_step, and so on, consumes exactly a gradient like this one) to actually update this layer's own learnable parameters.

From theory to code

Implement linear_backward(grad_output, x, weight), returning the three gradients (grad_x, grad_weight, grad_bias). grad_output has the same shape as the forward pass's output, (batch_size, out_features). Each of the three return values must have the SAME shape as the corresponding forward-pass input it's a gradient with respect to: grad_x matches x's shape, grad_weight matches weight's shape, grad_bias matches bias's shape.

Constraints

• grad_x must have the same shape as x: (batch_size, in_features).
• grad_weight must have the same shape as weight: (out_features, in_features).
• grad_bias must have the same shape as bias: (out_features,), summed (not averaged) across the batch dimension.
• Do not use any deep learning framework's autograd; derive and implement the gradients directly.

Hints