Where the Backward Pass Begins

A classifier sorts photos into $K$ classes. For one photo, its last layer emits $K$ raw scores, the logits $z_1, \dots, z_K$. Softmax turns them into probabilities, and categorical cross-entropy grades the probability given to the true class $y$:

Every backward pass through this network starts with the same question: how fast does the loss change when each logit is nudged? The chain from a logit $z_k$ to the loss $L$ has two links, and here are both of their rates.

• The loss link. The loss reads only the true class's probability, so $\dfrac{\partial L}{\partial p_y} = -\dfrac{1}{p_y}$. The other probabilities do not appear in $L$ at all.
• The softmax link. Nudging any logit moves $p_y$, because every logit sits in the shared denominator: $\dfrac{\partial p_y}{\partial z_k} = p_y\,\big([k = y] - p_k\big)$, where $[k = y]$ is $1$ when $k$ is the true class and $0$ otherwise.

A training batch holds $N$ photos, and the batch's loss is the mean of the $N$ per-photo losses.

Task: write logit_gradients(logits, labels).

• logits is a list of $N$ rows, each holding $K$ logits.
• labels is a list of $N$ class indices, counting from 0.

Return an $N \times K$ list of lists whose entry [i][k] is the rate of change of the batch loss with respect to logit $k$ of photo $i$, rounded to 4 decimal places.

A confident model can produce logits in the hundreds. Your function must still return correct, finite numbers for them.