Every activation earlier in this track (01-relu, 02-sigmoid, 03-tanh) is elementwise: each output value depends only on its own corresponding input value. Softmax breaks that pattern on purpose — a classifier's raw scores ("logits") only make sense as relative confidences, so turning them into a proper probability distribution requires every output to depend on every input, sharing one normalizing denominator across the whole row.
01-classical-ml/02-classification/06-softmax-cce's softmax already computed this forward pass. This question keeps that forward pass (normalizing over the last axis, like torch.softmax(x, dim=-1)) and adds softmax_backward(grad_output, output), which takes output (the saved forward result), same pattern as 02-sigmoid/03-tanh, and returns the vector-Jacobian product for the whole row at once — not a simple elementwise product, since softmax itself isn't elementwise.
x: shape (..., n_classes). softmax_forward returns the same shape, each row summing to 1.softmax_forward must stay finite for large |x| (no overflow from a raw exp).softmax_backward(grad_output, output) takes the forward pass's own saved result, not the original x, and returns a same-shape gradient computed per row (reducing across the last axis, not elementwise).Open one at a time. Each gives away a little more than the last.
Softmax isn't elementwise, so its backward pass can't be grad_output * f'(x) the way relu/sigmoid/tanh's are — every output in the row is coupled to every other output through the shared normalization.
The vector-Jacobian product collapses to output * (grad_output - dot), where dot is a single per-row scalar: the sum of grad_output * output across the row.
Click "Run Tests" to test your implementation