Numpy: get only elements on odd or even diagonal offsets in a matrix, change the rest to zeros
I have a 2D (square) matrix, for example it can be like this:
1 2 3
4 5 6
7 8 9

I want to get only the elements on the odd or even diagonal offsets of it, and let the rest be zeros. For example, with even diagonal offsets (%2 = 0), the resulting matrix is:
1 0 3
0 5 0
7 0 9

Explanation: main diagonal has offset 0, which is 1 5 9. The next diagonal offsets are 2, 6 and 4, 8 thus they are changed to zeros. Repeat the process until we reach the last diagonal.
And with odd diagonal index, the resulting matrix is:
0 2 0
4 0 6
0 8 0

I look at the np.diag(np.diag(x)) but it only returns the main diagonal and the rest are zeros. How can I extend it to odd/even offsets?
I can also use PyTorch.

I would do it following way using numpy
import numpy as np
arr = np.array([[1,2,3],[4,5,6],[7,8,9]])
masktile = np.array([[True,False],[False,True]])
mask = np.tile(masktile, (2,2))  # this must be at least as big as arr
arr0 = np.where(mask[:arr.shape[0],:arr.shape[1]], arr, 0)
arr1 = np.where(mask[:arr.shape[0],:arr.shape[1]], 0, arr)
print(arr0)
print(arr1)

output:
[[1 0 3]
 [0 5 0]
 [7 0 9]]
[[0 2 0]
 [4 0 6]
 [0 8 0]]

Explanation: I am creating mask which is array of Trues and Falses to use to decide if given element is to remain or should be replaced by 0. I create single tile which then I feed into np.tile to get &quot;chessboard&quot; of sufficient size, then I use part of apprioate size of it together with np.where to replace selected elements with 0.
