torch.diagonal¶
- torch.diagonal(input, offset=0, dim1=0, dim2=1) Tensor ¶
返回
input
的一个部分视图,其中相对于dim1
和dim2
的对角线元素作为维度添加到形状的末尾。参数
offset
控制考虑哪个对角线如果
offset
= 0,则是主对角线。如果
offset
> 0,则位于主对角线上方。如果
offset
< 0,则位于主对角线下方。
使用相同的参数对该函数的输出应用
torch.diag_embed()
会得到一个对角矩阵,其对角线元素与输入相同。但是,torch.diag_embed()
的默认维度不同,因此需要显式指定。- 参数
注意
要获取批次对角线,请传入 dim1=-2, dim2=-1。
示例
>>> a = torch.randn(3, 3) >>> a tensor([[-1.0854, 1.1431, -0.1752], [ 0.8536, -0.0905, 0.0360], [ 0.6927, -0.3735, -0.4945]]) >>> torch.diagonal(a) tensor([-1.0854, -0.0905, -0.4945]) >>> torch.diagonal(a, 1) tensor([ 1.1431, 0.0360]) >>> b = torch.randn(2, 5) >>> b tensor([[-1.7948, -1.2731, -0.3181, 2.0200, -1.6745], [ 1.8262, -1.5049, 0.4114, 1.0704, -1.2607]]) >>> torch.diagonal(b, 1, 1, 0) tensor([1.8262]) >>> x = torch.randn(2, 5, 4, 2) >>> torch.diagonal(x, offset=-1, dim1=1, dim2=2) tensor([[[-1.2631, 0.3755, -1.5977, -1.8172], [-1.1065, 1.0401, -0.2235, -0.7938]], [[-1.7325, -0.3081, 0.6166, 0.2335], [ 1.0500, 0.7336, -0.3836, -1.1015]]])