torch.diag_embed

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

创建一个张量，其中某些二维平面的对角线（由 dim1 和 dim2 指定）被 input 填充。为了便于创建批处理对角矩阵，默认选择返回张量的最后两个维度形成的二维平面。

offset 控制考虑哪个对角线：

• 如果 offset = 0，则是主对角线。

• 如果 offset > 0，则位于主对角线之上。

• 如果 offset < 0，则位于主对角线之下。

新矩阵的大小将计算为使指定对角线与最后一个输入维度的尺寸相匹配。请注意，对于除 $0$ 之外的其他 offset ， dim1 和 dim2 的顺序很重要。交换它们相当于改变 offset 的符号。

将 torch.diagonal() 应用于此函数的相同参数的输出，得到与输入相同的矩阵。然而， torch.diagonal() 具有不同的默认尺寸，因此需要显式指定这些尺寸。

参数:

• 输入（张量）- 输入张量。必须是至少 1 维的。

• 偏移量（int，可选）- 考虑哪个对角线。默认：0（主对角线）。

• dim1（int，可选）- 按照该维度取对角线。默认：-2。

• dim2（int，可选）- 按照该维度取对角线。默认：-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]]])