import numpy as np # create a 1D array with 1s in the diagonal arr = np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]) # reshape the array into a 3x3 matrix matrix = arr.reshape((3, 3)) print(matrix)