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() 具有不同的預設維度，因此需要顯式指定它們。

引數

• input (Tensor) – 輸入張量。必須至少是二維的。

• offset (int, optional) – 要考慮的對角線。預設值：0（主對角線）。

• dim1 (int, optional) – 用於提取對角線的第一個維度。預設值：0。

• dim2 (int, optional) – 用於提取對角線的第二個維度。預設值：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)
tensor([-1.0854, -0.0905, -0.4945])


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

>>> b = torch.randn(2, 5)
>>> b
tensor([[-1.7948, -1.2731, -0.3181,  2.0200, -1.6745],
        [ 1.8262, -1.5049,  0.4114,  1.0704, -1.2607]])

>>> torch.diagonal(b, 1, 1, 0)
tensor([1.8262])

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