python - 3D 数组乘法

Question

我有数组A和B两个维度 MxNxH。

我想定义一个二元运算符来“相乘”，这样结果就是 MxN 维度。

等效的操作是：

C = A[:,:,0] * B[:,:,0] + A[:,:,1] * B[:,:,1] + .... + A[:,:,H] * B[:,:,H]

有没有办法以更有效的方式执行此操作？
例如，在 numpy 中使用内置函数？

我试过tensordot了，但这给出了不同的结果。

Answer 1

最简单的是：

C = numpy.sum(A * B, -1)

我认为这也可能有效：

C = numpy.einsum("...i,...i->...", A, B)

Answer 2

试试这个：numpy.sum(A*B,axis=2)

这类似于其他建议，但可能更清楚（轴从 0 开始编号，因此轴 = 2 是第 3 轴或 MxNxH 中的 H）