python - 使用 numpy 使用 delta 函数来收缩索引

Question

我在 numpy 中有两个 NxM 数组，a并且b. 我想执行执行以下操作的矢量化操作：

c = np.zeros(N)
for i in range(N):
    for j in range(M):
        c[i] += a[i, j]*b[i, j]

用更数学的方式说，我有两个矩阵 A 和 B，并且想要计算矩阵 A*B 的对角线（矩阵转置等不精确）。我一直在尝试使用该tensordot功能完成类似的事情，但没有取得太大的成功。这是一个将要执行很多次的操作，所以我希望它是高效的（即，不从字面上计算矩阵 AB 并从中取对角线）。

标签： pythonnumpymatrix