torch.diag_embed¶
- torch.diag_embed(input, offset=0, dim1=-2, dim2=-1) Tensor ¶
创建一个张量,其某些 2D 平面(由
dim1
和dim2
指定)的对角线由input
填充。为了便于创建批处理对角矩阵,默认情况下会选择由返回张量的最后两个维度形成的 2D 平面。参数
offset
控制要考虑的对角线如果
offset
= 0,则为主对角线。如果
offset
> 0,则在主对角线之上。如果
offset
< 0,则在主对角线之下。
新矩阵的大小将被计算以使指定对角线的大小等于最后一个输入维度的尺寸。注意,对于
offset
不等于 的情况,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]]])