Managed to solve it via conv2d:
import torch
import torch.nn.functional as F

dim = 3
num_diagonls = 2*dim-1
# need to unsqueeze twice for use in conv2d
x = torch.rand(dim, dim).unsqueeze(0).unsqueeze(0)

print('x:')
print(x)

expected = torch.zeros(num_diagonls)
for i in range(num_diagonls):
    expected[i] = torch.sum(torch.diagonal(x[0][0], offset=dim-1-i))

print('expected diagonal sums:')
print(expected)

# need to unsqueeze twice for use in conv2d
w = torch.eye(dim).unsqueeze(0).unsqueeze(0)
# from after conv2d result, extract inner inner dim, then take the middle column
result = F.conv2d(x, w, padding=num_diagonls//2)[0][0][:, num_diagonls//2]
print('result diagonal sums:')
print(result)

outputs:
x:
tensor([[[[0.5106, 0.9985, 0.6385],
          [0.5529, 0.5057, 0.7685],
          [0.0656, 0.4694, 0.0604]]]])
expected diagonal sums:
tensor([0.6385, 1.7670, 1.0767, 1.0223, 0.0656])
result diagonal sums:
tensor([0.6385, 1.7670, 1.0767, 1.0223, 0.0656])

Roy