Normally we can sum over any dimention of a tensor, like this:
1+2+3
4+5+6
7+8+9
My question is how to sum all of the diagonals efficiently:
0+0+3
0+2+6
1+5+9
4+8+0
7+0+0
An equivalent question would be how to roll a tensor like this:
1 2 3
4 5 6
7 8 9
↓
1 2 3
5 6 4
9 7 8
And here is a naive solution:
a = torch.arange(1, 10).view(3,3)
b = torch.zeros_like(a)
for i in range(3)
 b[i] = a[i].roll(shifts=i, dims=0)

Obviously this will be slow because of the for loop, how should I get rid of the for loop?