利用到Numpy库函数 numpy.diag_indices_from

import numpy as np

#3×3的单位矩阵
a = np.eye(3)                

#获取主对角线元素的索引
row, col = np.diag_indices_from(a) 

#对主对角线进行赋值
a[row, col] = [3, 3, 3]
#or
a[row, col] = 3

print(a)

#比较实用的是将某一个计算结果依次赋给主对角线元素
例如
a[row, col] = np.array([sum(w) for w in W])
输出:

[[3, 0, 0]
 [0, 3, 0]
 [0, 0, 3]]