Issue
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
?
Solution
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]])
Answered By - llllllllll
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.