Elegant Numpy Tensor product

Question

I need to take the product over two tensors in numpy (or pytorch):

I have

A = np.arange(1024).reshape(8,1,128)
B = np.arange(9216).reshape(8, 128, 9)

And want to obtain C, with dot products summing over the last dim of A (axis=2) and the middle dim of B (axis=1). This should have dimensions 8x9. Currently, I am doing:

C = np.zeros([8, 9])
for i in range(8):
    C[i,:] = np.matmul(A[i,:,:], B[i,:,:])

How to do this elegantly?

I tried:

np.tensordot(weights, features, axes=(2,1)).

but it returns 8x1x8x9.

Answer 1

One way would be to use numpy.einsum.

C = np.einsum('ijk,ikl->il', A, B)

Or you could use broadcasted matrix multiply.

C = (A @ B).squeeze(axis=1)
# equivalent: C = np.matmul(A, B).squeeze(axis=1)