torch.diagonal#

torch.diagonal(input,offset=0,dim1=0,dim2=1)→Tensor#

Returns a partial view ofinput with the its diagonal elementswith respect todim1 anddim2 appended as a dimensionat the end of the shape.

The argumentoffset controls which diagonal to consider:

• Ifoffset = 0, it is the main diagonal.

• Ifoffset > 0, it is above the main diagonal.

• Ifoffset < 0, it is below the main diagonal.

Applyingtorch.diag_embed() to the output of this function withthe same arguments yields a diagonal matrix with the diagonal entriesof the input. However,torch.diag_embed() has different defaultdimensions, so those need to be explicitly specified.

Parameters

• input (Tensor) – the input tensor. Must be at least 2-dimensional.

• offset (int,optional) – which diagonal to consider. Default: 0(main diagonal).

• dim1 (int,optional) – first dimension with respect to which totake diagonal. Default: 0.

• dim2 (int,optional) – second dimension with respect to which totake diagonal. Default: 1.

Note

To take a batch diagonal, pass in dim1=-2, dim2=-1.

Examples:

>>>a=torch.randn(3,3)>>>atensor([[-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)>>>btensor([[-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]]])