小编non*_*ono的帖子

我有一个形状为 ( m*n , m*n ) 的张量，我想提取一个大小为 ( n , m*n )的张量，其中包含对角线上大小为n*n的 m 个块。例如：

>>> a
tensor([[1, 2, 0, 0],
        [3, 4, 0, 0],
        [0, 0, 5, 6],
        [0, 0, 7, 8]])

我想要一个extract(a, m, n)将输出的函数：

>>> extract(a, 2, 2)
tensor([[1, 2, 5, 6],
        [3, 4, 7, 8]])

我想过使用某种切片，因为块可以表示为：

>>> for i in range(m):
...     print(a[i*m: i*m + n, i*m: i*m + n])
tensor([[1, 2],
        [3, 4]])
tensor([[5, 6],
        [7, 8]])