pytorch 变对角矩阵_在numpy或pytorch中自动获得对角矩阵条纹

我需要获得矩阵的对角条纹(这里不确定术语，对角矩阵条纹似乎最好描述它) .

比方说，我有一个大小为KxN的矩阵，其中K和N是任意大小，K> N.说，我有一个矩阵：

[[ 0 1 2]

[ 3 4 5]

[ 6 7 8]

[ 9 10 11]]

从中我需要提取对角条纹，在这种情况下，是通过截断原始条带创建的矩阵MxV大小：

[[ 0 x x]

[ 3 4 x]

[ x 7 8]

[ x x 11]]

所以结果矩阵是：

[[ 0 4 8]

[ 3 7 11]]

下面是一个使用矩阵屏蔽的小示例代码，用于去除屏蔽位置：

import numpy as np

X=np.arange(12).reshape(4,3)

mask=np.asarray([

[ True, False, False],

[ True, True, False],

[ False, True, True],

[ False, False, True]

])

>>> mask

array([[ True, False, False],

[ True, True, False],

[False, True, True],

[False, False, True]], dtype=bool)

>>> X

array([[ 0, 1, 2],

[ 3, 4, 5],

[ 6, 7, 8],

[ 9, 10, 11]])

>>> X.T[mask.T].reshape(3,2).T

array([[ 0, 4, 8],

[ 3, 7, 11]])

但我不知道如何将这种面具自动生成任何K和N尺寸，例如 . 39x9或360x96

任何帮助表示赞赏 . 也许有一些函数可以在numpy，scipy或pytorch中自动执行此操作？

编辑：

我有另一个问题，是否有可能而不是得到：

[[ 0 x x]

[ 3 4 x]

[ x 7 8]

[ x x 11]]

要得到这样的反向条纹：

[[ x x 2]

[ x 4 5]

[ 6 7 x]

[ 9 x x]]