megengine.functional.diag¶

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

Extract a diagonal or construct a diagonal tensor.

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

参数

inp – 输入张量。
k (int) – diagonal in consider. Use \(k=0\) for the main diagonal, \(k>0\) for diagonals above the main diagonal, and \(k<0\) for diagonals below the main diagonal.

参见

If you want to create a identity matrix, see eye.

返回类型: Tensor
返回: the extracted diagonal or constructed diagonal tensor.

实际案例

Input is a 1D 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)

Input is a 2D tensor:

>>> 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)

Get the k-th diagonal of a given matrix:

>>> 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