3D 数组乘法

Question

I have arrays A and B both of dimension MxNxH.

I would like to define a binary operator, to "multiply", such that the result is MxN dimensions.

The equivalent operation would be:

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

Is there a way to do this operation in a more efficient way?
For example, using a built in function in numpy?

I have tried tensordot , but this gives a different result.

Answer 1

The easiest is:

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

I think this might work too:

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

Answer 2

try this: numpy.sum( A*B, axis=2 )

this is similar to the other suggestion but perhaps clearer (axes are numbered from 0, so axis=2 is the 3rd axis or H out of MxNxH)