@Dynasty thanks for posting those formulas, that’s quite interesting. The downside is that they always end with a dot product with this a vector, where you substitute it with a one-hot value to get a single entry. Note that removing the dot product would just give you a row of the hessian. In torch.autograd (or any other reverse-mode AD software that doesn’t do partial evaluation - I’m not aware of any) this would be equivalent to just computing the full row, and then only taking the diagonal entry, so there are no savings