Elegant Numpy Tensor product
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.

One way would be to use numpy.einsum.
C = np.einsum('ijk,ikl-&gt;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)

