如何在numpy中得到nXn矩阵的向量的向量对角线

Question

例如，我有

import numpy as np
a = np.ones((100, 5, 5))

我想要

d = np.vector_diagonal(a)
assert d.shape == (100, 5)

其中d[i, j]对应于a[i, j, j]

如何使用numpy？

Answer 1

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

只需要选择哪些轴是"；矩阵"；并且其中"；矢量化矩阵"；

将在选定的轴上进行缩减。