How to apply the gradient of softmax in backprop
After further working on this, I figured out that:


The homework implementation combines softmax with cross entropy loss as a matter of choice, while my choice of keeping softmax separate as an activation function is also valid.
The homework implementation is indeed missing the derivative of softmax for the backprop pass.
The gradient of softmax with respect to its inputs is really the partial of each output with respect to each input:




So for the vector (gradient) form:


Which in my vectorized numpy code is simply:

self.data * (1. - self.data)


Where self.data is the softmax of the input, previously computed from the forward pass.