torch.diag_embed¶

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

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

参数 offset 控制要考虑的对角线：

如果 offset = 0，则是主对角线。
如果 offset 大于 0，则位于主对角线上方。
如果 offset 小于 0，则位于主对角线下方。

新矩阵的大小将根据最后一个输入维度的大小计算，以确保指定对角线的位置。请注意，对于offset（偏移量）不为 $0$ 的情况，dim1 和dim2的顺序很重要。交换它们的顺序相当于改变offset的符号。

将torch.diagonal()应用于该函数的输出，并使用相同的参数，将得到与输入相同的矩阵。然而，torch.diagonal()具有不同的默认维度，因此需要明确指定这些维度。

Parameters

输入 (张量) – 输入张量。必须至少是1维的。
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]]])