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?

You can 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]])

