How to get a tensor of diagonal vectors

I have a tensor which has some diagonal matrices inside it.

tensor([[[-0.2920,  0.0000],
         [ 0.0000, -0.8663]],

        [[-0.2920,  0.0000],
         [ 0.0000, -0.7568]],

        [[-0.3373,  0.0000],
         [ 0.0000, -0.5953]],

        [[-0.0068,  0.0000],
         [ 0.0000, -0.8065]],

        [[-0.1260,  0.0000],
         [ 0.0000, -0.7147]]], grad_fn=<MulBackward0>)

I want to do an operation on this to change this tensor to the following such that I do not lose the gradients because I need to call backward later.

tensor([[-0.2920, -0.8663],

        [-0.2920, -0.7568],

        [-0.3373, -0.5953],

        [-0.0068, -0.8065],

        [-0.1260, -0.7147]])

I saw diagonal_embed but it didnt help much. How do I do this?

Hi @Ananyapam_De,

You should be able to do something like,

diagonal_tensors = torch.randn(10, 2).diag_embed() #creates batch of diagonal matrices
diagonal_only = torch.diagonal(diagonal_tensors, 0, -2, -1) #get 0 = diagonal and -2,-1 is over the last two dimensions. 
#returns 
"""
tensor([[ 1.4204,  0.3541],
        [-0.0236, -0.3348],
        [ 0.0364,  0.7933],
        [ 0.4506,  1.5018],
        [-1.2026, -0.0438],
        [ 1.8065,  1.8191],
        [-1.2577,  0.7627],
        [ 0.6864, -0.3432],
        [-0.0303,  0.0384],
        [ 0.0991,  0.4471]])
"""

if you want it the same shape as your example (with the spacing between the rows), you’ll have to unsqueeze the results, i.e.,

 torch.diagonal(diagonal_tensors, 0, -2, -1).unsqueeze(-2)
"""
#returns 
tensor([[[ 1.4204,  0.3541]],

        [[-0.0236, -0.3348]],

        [[ 0.0364,  0.7933]],

        [[ 0.4506,  1.5018]],

        [[-1.2026, -0.0438]],

        [[ 1.8065,  1.8191]],

        [[-1.2577,  0.7627]],

        [[ 0.6864, -0.3432]],

        [[-0.0303,  0.0384]],

        [[ 0.0991,  0.4471]]])
"""

But you can just check the shapes via x.shape if you need to.