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()具有不同的默认维度，因此需要显式指定这些维度。

参数

• 输入 (Tensor) – 必须是一个至少一维的输入张量。

• offset (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]]])