In [219]: m = mat.shape[-1]
In [220]: mat2 = np.zeros(mat.shape+(m,),mat.dtype)
In [221]: idx = np.arange(m)
In [222]: mat2[...,idx,idx] = mat[...,idx]
In [223]: mat2
Out[223]:
array([[[1, 0],
[0, 1]],
[[2, 0],
[0, 2]],
[[3, 0],
[0, 3]]])
тестирование на большем mat
:
In [224]: mat = np.array([[1,2,3],[2,3,4],[3,4,5]])
In [225]: m = mat.shape[-1]
In [226]: mat2 = np.zeros(mat.shape+(m,),mat.dtype)
In [227]: idx = np.arange(m)
In [228]: mat2[...,idx,idx] = mat[...,idx]
In [230]: mat2
Out[230]:
array([[[1, 0, 0],
[0, 2, 0],
[0, 0, 3]],
[[2, 0, 0],
[0, 3, 0],
[0, 0, 4]],
[[3, 0, 0],
[0, 4, 0],
[0, 0, 5]]])