numpy.diag(v, k=0)
can extract the diagonal of an array or can construct a diagonal array.
If v is a 2-D array, it returns a copy of its k-th diagonal.
>>> x = np.array([[11, 12, 13], [21, 22, 23], [31, 32, 33]])
>>> np.diag(x)
array([11, 22, 33])
>>> np.diag(x, k=1)
array([12, 23])
>>> np.diag(x, k=-1)
array([21, 32])
If v is a 1-D array, it returns a 2-D array with v on the k-th diagonal.
>>> np.diag([3,4,5])
array([[3, 0, 0],
[0, 4, 0],
[0, 0, 5]])
>>> np.diag([3,4,5], k=1)
array([[0, 3, 0, 0],
[0, 0, 4, 0],
[0, 0, 0, 5],
[0, 0, 0, 0]])