numpy.identity(n, dtype=None)
Returns a square array (n x n) with ones on the main diagonal.
>>> import numpy as np
>>> np.identity(3, int)
array([[1, 0, 0],
[0, 1, 0],
[0, 0, 1]])
>>>
>>> # default is float
>>> np.identity(3)
array([[ 1., 0., 0.],
[ 0., 1., 0.],
[ 0., 0., 1.]])