torch.diag

torch.diag(input, diagonal=0, *, out=None) → Tensor

如果 input 是一个向量（即一维张量），则返回一个二维的方形张量，其对角线上的元素与 input 中的元素相同。
如果 input 是一个矩阵（二维张量），则返回其对角线元素构成的一维张量。

参数 diagonal 控制要考虑哪一条对角线：

当 diagonal = 0 时，表示的是主对角线。
如果 diagonal > 0，表示它位于主对角线的上方。
如果 diagonal 小于 0，它位于主对角线的下方。

参数

input (Tensor) – 需要输入的张量。
diagonal (int, 可选) – 需要考虑的对角线

关键字参数

out (Tensor, 可选) – 指定输出张量。

参见

torch.diagonal() 总是返回输入张量的对角线元素。

torch.diagflat() 总是根据输入创建一个具有指定对角元素的张量。

示例:

生成一个以输入向量为对角线元素的方阵。

>>> a = torch.randn(3)
>>> a
tensor([ 0.5950,-0.0872, 2.3298])
>>> torch.diag(a)
tensor([[ 0.5950, 0.0000, 0.0000],
        [ 0.0000,-0.0872, 0.0000],
        [ 0.0000, 0.0000, 2.3298]])
>>> torch.diag(a, 1)
tensor([[ 0.0000, 0.5950, 0.0000, 0.0000],
        [ 0.0000, 0.0000,-0.0872, 0.0000],
        [ 0.0000, 0.0000, 0.0000, 2.3298],
        [ 0.0000, 0.0000, 0.0000, 0.0000]])

获取矩阵的第 k 条对角线：

>>> a = torch.randn(3, 3)
>>> a
tensor([[-0.4264, 0.0255,-0.1064],
        [ 0.8795,-0.2429, 0.1374],
        [ 0.1029,-0.6482,-1.6300]])
>>> torch.diag(a, 0)
tensor([-0.4264,-0.2429,-1.6300])
>>> torch.diag(a, 1)
tensor([ 0.0255, 0.1374])