Starting from a 2d identity matrix, here are two options you can make the "3d identity matrix":
import numpy as np
i = np.identity(2)
Option 1: stack the 2d identity matrix along the third dimension
np.dstack([i]*3)
#array([[[ 1., 1., 1.],
# [ 0., 0., 0.]],
# [[ 0., 0., 0.],
# [ 1., 1., 1.]]])
Option 2: repeat values and then reshape
np.repeat(i, 3, axis=1).reshape((2,2,3))
#array([[[ 1., 1., 1.],
# [ 0., 0., 0.]],
# [[ 0., 0., 0.],
# [ 1., 1., 1.]]])
Option 3: Create an array of zeros and assign 1 to positions of diagonal elements (of the 1st and 2nd dimensions) using advanced indexing:
shape = (2,2,3)
identity_3d = np.zeros(shape)
idx = np.arange(shape[0])
identity_3d[idx, idx, :] = 1
identity_3d
#array([[[ 1., 1., 1.],
# [ 0., 0., 0.]],
# [[ 0., 0., 0.],
# [ 1., 1., 1.]]])
Timing:
%%timeit
shape = (100,100,300)
i = np.identity(shape[0])
np.repeat(i, shape[2], axis=1).reshape(shape)
# 10 loops, best of 3: 10.1 ms per loop
%%timeit
shape = (100,100,300)
i = np.identity(shape[0])
np.dstack([i] * shape[2])
# 10 loops, best of 3: 47.2 ms per loop
%%timeit
shape = (100,100,300)
identity_3d = np.zeros(shape)
idx = np.arange(shape[0])
identity_3d[idx, idx, :] = 1
# 100 loops, best of 3: 6.31 ms per loop