PyTorch torch.diag_embed 函数

torch.diag_embed 是 PyTorch 中用于将输入张量作为对角线元素嵌入到新张量中的函数。它在输入的最后一个维度上创建对角线。

torch.diag_embed(input, diagonal=0, offset=0)

参数说明：

使用示例

import torch

# 创建一维张量
x = torch.tensor([1, 2, 3])

# 创建对角嵌入
y = torch.diag_embed(x)
print(y)

输出结果为：

tensor([[1, 0, 0],
        [0, 2, 0],
        [0, 0, 3]])

import torch

# 创建二维张量
x = torch.tensor([[1, 2], [3, 4]])

# 创建对角嵌入
y = torch.diag_embed(x)
print(y.shape)
print(y)

输出结果为：

torch.Size([2, 2, 2])
tensor([[[1, 0],
         [0, 2]],

        [[3, 0],
         [0, 4]]])