Extracting the diagonal elements of a square matrix

For square matrices, we can extract the diagonal elements

>>> C = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> print(C)
[[1 2 3]
 [4 5 6]
 [7 8 9]]
>>> print(np.diag(C))
[1 5 9]