获取非对称张量的对角线以完美地结束在 pytorch 的右下角

Question

我有一个形状为

(5 * n, n)

的张量，我基本上想从第一列的前 5 行中提取前 5 个元素，然后移过 1 列并为第二列提取接下来的 5 行，然后移过等等. 它有点像 torch.diagonal 但适用于非对称张量，并且可以假设张量具有适当的尺寸，因此每次都能正常工作。

例如，如果我的输入张量是这样的：

>>> t = torch.arange(45).reshape(15, 3)
>>> t
tensor([[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17],
        [18, 19, 20],
        [21, 22, 23],
        [24, 25, 26],
        [27, 28, 29],
        [30, 31, 32],
        [33, 34, 35],
        [36, 37, 38],
        [39, 40, 41],
        [42, 43, 44]])

然后我想要一些方法来获得

out = tensor([0, 3, 6, 9, 12, 16, 19, 22, 25, 28, 32, 35, 38, 41, 44])

我更喜欢我不必使用循环，因为我会为一些相当大的输入做上千次，所以我觉得那样效率很低。