torch.diagonal

torch.diagonal(input, offset=0, dim1=0, dim2=1) → Tensor

返回input的部分视图，并将相对于dim1和dim2的对角线元素作为形状末尾的一个新维度添加进去。

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

当 offset = 0 时，表示为主对角线。
如果 offset 大于 0，它位于主对角线的上方。
如果 offset 小于 0，它位于主对角线之下。

使用torch.diag_embed()函数对本函数的输出进行处理，并保持参数相同，可以得到一个以输入矩阵对角线元素为对角线的新对角矩阵。不过，torch.diag_embed()具有不同的默认维度设置，因此需要显式指定这些维度。

参数

输入 (Tensor) – 必须是一个至少二维的输入张量。
offset (int, 可选) – 指定要考虑的对角线。默认值：0（主对角线）。
dim1 (int, 可选) – 第一个维度，用于取对角线。默认值：0。
dim2 (int, 可选) – 相对于该维度取对角线的第二个维度。默认值：1。

注意

要获取一批对角元素，请将 dim1 设为 -2，dim2 设为 -1。

示例:

>>> a = torch.randn(3, 3)
>>> a
tensor([[-1.0854,  1.1431, -0.1752],
        [ 0.8536, -0.0905,  0.0360],
        [ 0.6927, -0.3735, -0.4945]])


>>> torch.diagonal(a, 0)
tensor([-1.0854, -0.0905, -0.4945])


>>> torch.diagonal(a, 1)
tensor([ 1.1431,  0.0360])


>>> x = torch.randn(2, 5, 4, 2)
>>> torch.diagonal(x, offset=-1, dim1=1, dim2=2)
tensor([[[-1.2631,  0.3755, -1.5977, -1.8172],
         [-1.1065,  1.0401, -0.2235, -0.7938]],

        [[-1.7325, -0.3081,  0.6166,  0.2335],
         [ 1.0500,  0.7336, -0.3836, -1.1015]]])