CODEWITHSUNDEEP
pythonnumpypytorch
Ujan
Ujan

Reputation: 103

Get all 2D diagonals of a 3D tensor in numpy

I have a 3D tensor A x B x C. For each matrix B x C, I want to extract the leading diagonal.

Is there a vectorized way of doing this in numpy or pytorch instead of looping over A?

Upvotes: 2

Views: 1022

Answers (3)

llllllllll
llllllllll

Reputation: 16404

Use numpy.diagonal():

np.diagonal(a, axis1=1, axis2=2)

Example:

In [10]: a = np.arange(3*4*5).reshape(3,4,5)

In [11]: a
Out[11]: 
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]],

       [[20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39]],

       [[40, 41, 42, 43, 44],
        [45, 46, 47, 48, 49],
        [50, 51, 52, 53, 54],
        [55, 56, 57, 58, 59]]])

In [12]: np.diagonal(a, axis1=1, axis2=2)
Out[12]: 
array([[ 0,  6, 12, 18],
       [20, 26, 32, 38],
       [40, 46, 52, 58]])

Upvotes: 3

iacob
iacob

Reputation: 24231

In PyTorch, use torch.diagonal():

t.diagonal(dim0=-2, dim1=-1)

Upvotes: 0

Divakar
Divakar

Reputation: 221584

Assuming that the leading diagonal for a generic non-squared (BxC) slice starts off from the top-left corner, we can reshape and slice -

a.reshape(a.shape[0],-1)[:,::a.shape[-1]+1]

Sample run -

In [193]: np.random.seed(0)

In [194]: a = np.random.randint(11,99,(3,4,5))

In [195]: a
Out[195]: 
array([[[55, 58, 75, 78, 78],
        [20, 94, 32, 47, 98],
        [81, 23, 69, 76, 50],
        [98, 57, 92, 48, 36]],

       [[88, 83, 20, 31, 91],
        [80, 90, 58, 75, 93],
        [60, 40, 30, 30, 25],
        [50, 43, 76, 20, 68]],

       [[43, 42, 85, 34, 46],
        [86, 66, 39, 45, 11],
        [11, 47, 64, 16, 49],
        [28, 90, 15, 53, 69]]])

In [196]: a.reshape(a.shape[0],-1)[:,::a.shape[-1]+1]
Out[196]: 
array([[55, 94, 69, 48],
       [88, 90, 30, 20],
       [43, 66, 64, 53]])

Upvotes: 1

Related Questions