mlx.core.diag

目录

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