megengine.functional.diag¶

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

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

参数

inp – 输入张量。
k – 对角线偏移量。使用 \(k=0\) 则表示在主对角线上， \(k>0\) 表示在主对角线上方， \(k<0\) 表示在主对角线下方。默认：0.

返回类型

Tensor

返回

提取出的对角线或构造出的对角矩阵。

实际案例

>>> inp = F.arange(6, dtype='int32').reshape(2,3)
>>> out = F.diag(inp, k=1)
>>> out
Tensor([1 5], dtype=int32, device=xpux:0)
>>> F.diag(out)
Tensor([[1 0]
 [0 5]], dtype=int32, device=xpux:0)

megengine.functional.matmul

megengine.functional.svd