torch.diag_embed

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

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

参数 offset 控制考虑哪个对角线

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

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

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

新矩阵的大小将计算出来，以使指定的对角线尺寸与最后一个输入维度的大小一致。请注意，对于 offset 不等于 $0$ 的情况，dim1 和 dim2 的顺序很重要。交换它们等同于改变 offset 的符号。

使用相同的参数对该函数的输出应用 torch.diagonal() 会产生一个与输入相同的矩阵。然而，torch.diagonal() 有不同的默认维度，因此需要明确指定。

参数

• input (Tensor) – 输入 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]]])