torch.diag

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

• 如果 input 是一個向量（1 維張量），則返回一個 2 維正方形張量，其中 input 的元素作為對角線。

• 如果 input 是一個矩陣（2 維張量），則返回一個 1 維張量，其中包含 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])