numpy.diag

numpy.diag 是 NumPy 库中的一个函数，用于提取一个矩阵的对角线元素或者构造一个对角线矩阵。

具体来说，当输入一个矩阵时，函数会返回一个由原矩阵的对角线元素构成的一维数组；当输入一个一维数组时，函数会返回一个以该数组为对角线元素的方阵。

例如，对于以下矩阵：

[[1, 2, 3],
 [4, 5, 6],
 [7, 8, 9]]

调用 `np.diag` 函数可以得到该矩阵的对角线元素 `[1, 5, 9]`；调用 `np.diag([1, 2, 3])` 函数可以得到以 `[1, 2, 3]` 为对角线元素的方阵：

[[1, 0, 0],
 [0, 2, 0],
 [0, 0, 3]]