mlx.core.diag# diag(a: array, /, k: int = 0, *, stream: None | Stream | Device = None) → array# 提取对角线或构建对角矩阵。如果 a 是 1 维,则构建一个以 a 作为第 \(k\) 条对角线的对角矩阵。如果 a 是 2 维,则返回第 \(k\) 条对角线。 参数: a (array) – 1 维或 2 维输入数组。 k (int, 可选) – 要提取或构建的对角线。默认值:0。 返回值: 提取的对角线或构建的对角矩阵。 返回类型: array