Necesito obtener una franja diagonal de la matriz (no estoy seguro de la terminología aquí, la franja de matriz diagonal parece describirlo mejor).

Digamos que tengo una matriz de tamaño KxN, donde K y N son tamaños arbitrarios y K> N. Digamos que tengo una matriz:

[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]

De él necesitaría extraer una franja diagonal, en este caso, un tamaño de matriz MxV que se crea truncando el original:

[[ 0  x  x]
 [ 3  4  x]
 [ x  7  8]
 [ x  x  11]]

Entonces la matriz de resultados es:

[[ 0  4  8]
 [ 3  7  11]]

Aquí hay un pequeño código de ejemplo que usa el enmascaramiento para las matrices, para quitar las posiciones enmascaradas:

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]])

Pero no veo cómo una máscara de este tipo podría generarse automáticamente a cualquier tamaño K y N, por ejemplo. 39x9 o 360x96

Cualquier ayuda es apreciada. ¿Quizás hay alguna función que hace esto automáticamente en numpy, scipy o pytorch?

EDITAR:

Tengo otra pregunta, ¿es posible en lugar de obtener:

[[ 0  x  x]
 [ 3  4  x]
 [ x  7  8]
 [ x  x  11]]

Para obtener una franja inversa como esta:

[[ x   x   2]
 [ x   4   5]
 [ 6   7   x]
 [ 9   x   x]]