Python PyTorch diag_embed用法及代码示例

本文简要介绍python语言中 torch.diag_embed 的用法。

用法:

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

参数：

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

• offset(int,可选的) -考虑哪个对角线。默认值：0(主对角线)。

• dim1(int,可选的) -对角线的第一个维度。默认值：-2。

• dim2(int,可选的) -相对于其取对角线的第二维。默认值：-1。

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

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

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

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

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

新矩阵的大小将被计算为最后一个输入维度的大小的指定对角线。请注意，对于 offset 以外的 0 ， dim1 和 dim2 的顺序很重要。交换它们相当于更改 offset 的符号。

使用相同的参数将 torch.diagonal() 应用于此函数的输出会产生与输入相同的矩阵。但是， torch.diagonal() 具有不同的默认尺寸，因此需要明确指定。

例子：

>>> 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]]])