check with Numpy.

import numpy as np
def idmatrix_np(size):
    return np.identity(size)

Example or Call Funtion.

size = 4 
identity_matrix = idmatrix_np(size)