I am attempting to re-implement backpropagation on my own for didactic purposes, but am running into some issues. I am trying to work backwards from a simple network, starting with LogSoftmax + NLLLoss, but I am unable to match the calculated gradient of the input to the LogSoftmax layer as calculated by autograd.
import torch
import torch.nn

new_relu_feats = torch.Tensor([[1,0,3]])
new_relu_feats.requires_grad = True
logits = torch.nn.LogSoftmax(dim=1)(new_relu_feats)
logits.retain_grad()
label = torch.LongTensor([1])
loss = torch.nn.NLLLoss()(logits, label)
loss.backward()

sm = torch.nn.Softmax(dim=1)(new_relu_feats)
dloss_dlogits = torch.Tensor([[0,-1,0]])
logits.grad # This matches above
dlogits_dnew_relu_feat = torch.Tensor([[[1-sm[0][0], -sm[0][0], -sm[0][0]], [-sm[0][1], 1-sm[0][1], -sm[0][1]], [-sm[0][2], -sm[0][2], 1-sm[0][2]]]])

dloss_dlogits * dlogits_dnew_relu_feat
# This does not entirely match above, but the middle column does (corresponding to the correct class)
new_relu_feats.grad

Is this correct, but the matrix is eventually reshaped for efficiency by simply selecting the portions that are non-zero (aka corresponding to the correct class)?