你可以使用einsum来提取对角线。你会得到一个原始数组的视图，这意味着如果你修改视图，你也会修改原始数组。这可以让你覆盖原始数组的条目，就像这样：

# setup
import numpy as np
X = np.arange(0, 3*4*5).reshape((3, 4, 5))  # bogus data
Y = np.zeros((3, 4, 5, 5))                  # initialize Y with zeros

# extract diagonal, and write into it
Y_diag_view = np.einsum('ijkk->ijk', Y)
Y_diag_view[...] = X                        # modify Y_diag_view and therefore X in place
print(Y[0, 0, ...])                         # prints diagonal matrix