三維數組乘法

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，軸= 2）

這與其他建議類似，但可能更清晰（軸從0開始編號，因此axis = 2是第3軸或H out of MxNxH）