torch.diag_embed

torch.diag_embed(input, offset=0, dim1=-2, dim2=-1) → Tensor

創建一個張量，其特定 2D 平面（由 dim1 和 dim2 指定）的對角線會被 input 填充。為了方便創建批次對角矩陣，預設會選擇返回張量的最後兩個維度形成的 2D 平面。

參數 offset 控制要考慮哪個對角線

• 如果 offset = 0，則為主對角線。

• 如果 offset > 0，則在主對角線上方。

• 如果 offset < 0，則在主對角線下方。

新矩陣的大小將被計算，以使指定對角線的大小與最後一個輸入維度的大小相同。請注意，對於 offset 不是 $0$ 的情況，dim1 和 dim2 的順序很重要。交換它們等同於更改 offset 的符號。

使用相同的參數將 torch.diagonal() 應用於此函數的輸出會產生與輸入相同的矩陣。但是，torch.diagonal() 具有不同的預設維度，因此需要明確指定這些維度。

參數

• input (Tensor) – 輸入張量。必須至少為 1 維。

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

• dim1 (int, optional) – 要對其取對角線的第一個維度。預設值：-2。

• dim2 (int, optional) – 要對其取對角線的第二個維度。預設值：-1。

範例

>>> a = torch.randn(2, 3)
>>> torch.diag_embed(a)
tensor([[[ 1.5410,  0.0000,  0.0000],
         [ 0.0000, -0.2934,  0.0000],
         [ 0.0000,  0.0000, -2.1788]],

        [[ 0.5684,  0.0000,  0.0000],
         [ 0.0000, -1.0845,  0.0000],
         [ 0.0000,  0.0000, -1.3986]]])

>>> torch.diag_embed(a, offset=1, dim1=0, dim2=2)
tensor([[[ 0.0000,  1.5410,  0.0000,  0.0000],
         [ 0.0000,  0.5684,  0.0000,  0.0000]],

        [[ 0.0000,  0.0000, -0.2934,  0.0000],
         [ 0.0000,  0.0000, -1.0845,  0.0000]],

        [[ 0.0000,  0.0000,  0.0000, -2.1788],
         [ 0.0000,  0.0000,  0.0000, -1.3986]],

        [[ 0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000]]])