Some time ago, GitHub - cybertronai/autograd-lib 1 had something for this.
The new way is supposed to be with vmap, but it seems it is not quite there yet [feature request] Simple and Efficient way to get gradients of each element of a sum · Issue #7786 · pytorch/pytorch · GitHub .
I did it manually when I needed it.
I also want to write a source-to-source differentiation thing for TorchScript if I ever find the time.
Best regards
Thomas