megengine.functional.diag¶

diag(inp, k=0)[源代码]¶

提取或者构造一个对角线向量

如果 inp 是 1 维张量，则返回由 inp 中元素作为对角线的 2 维张量。如果 inp 是 2 维张量，则返回 inp 对角线元素构成的 1 维张量。

参数

inp – 输入张量。
k (int) – 控制对角线偏移量。如果 \(k=0\) 则为下对角线。

参见

如果像创建个单位矩阵，请见 eye

返回类型: Tensor
返回: 提取或者构造对角向量。

实际案例

输入是一个一维向量

>>> F.diag(Tensor([1, 2, 3]))
Tensor([[1 0 0]
 [0 2 0]
 [0 0 3]], dtype=int32, device=xpux:0)
>>> F.diag(Tensor([1, 2, 3]), k=1)
Tensor([[0 1 0 0]
 [0 0 2 0]
 [0 0 0 3]
 [0 0 0 0]], dtype=int32, device=xpux:0)

输入是一个二维向量：

>>> x = F.arange(9).reshape(3, 3)
>>> x
Tensor([[0. 1. 2.]
 [3. 4. 5.]
 [6. 7. 8.]], device=xpux:0)
>>> F.diag(x)
Tensor([0. 4. 8.], device=xpux:0)

求给定矩阵的第k个对角线：

>>> F.diag(x, k=1)
Tensor([1. 5.], device=xpux:0)
>>> F.diag(x, k=-1)
Tensor([3. 7.], device=xpux:0)

megengine.functional.eye

megengine.functional.zeros