multiplying a vector (1 x N) by a tensor (N x M x M)
I am looking for a matrix operation in numpy or preferably in pytorch that allows one to multiply a vector (1 x N) by a tensor (N x M x M) and get (1 x M x M).  This is easily accomplished using a for loop, but the for loop does not allow back propagation during training. I tried using matmul in numpy and pytorch (and several others such as dot and bmm), but could not get any to work. Here is an example (where M=2, but is 256 in my use case) of what I am trying to do:
a = np.array([1,2,3])
b = np.array([[[1,2],[3,4]],[[5,6],[7,8]],[[9,10],[11,12]]])

I would like to perform the operation: 1*[[1,2],[3,4]] + 2*[[5,6],[7,8]] + 3*[[9,10],[11,12]], which can be achieved with a for loop like this:
for i in range(3):
  matrix_sum += a[i]*b[i]

Any advice or solution would be greatly appreciated.

Numpy and pytorch were built uppon matrix multiplications!
Torch example:
A = torch.rand(1, N)
B = torch.rand(N, M, M)

C = A @ B.transpose(0, 1) 
C.transpose_(0, 1)
C.shape
torch.size(1, M, M)

And similarly for numpy:
A = np.random.randn(1, N)
B = np.random.randn(N, M, M)

C = A @ B.transpose(1, 0, 2)
C = C.transpose(1, 0, 2)
C.shape
(1, M, M)

Edit For the Einsum lovers:
Pytorch and numpy handle einsum pretty much in the same way:
torch.einsum('i,ijk-&gt;jk', A, B)
np.einsum('i,ijk-&gt;jk', A, B)

Pytorch einsum documentation: https://pytorch.org/docs/master/generated/torch.einsum.html
Numpy einsum documentation: https://numpy.org/doc/stable/reference/generated/numpy.einsum.html
