Pytorch batch matrix vector outer product

Question

I am trying to generate a vector-matrix outer product (tensor) using PyTorch. Assuming the vector v has size p and the matrix M has size qXr, the result of the product should be pXqXr.

Example:

#size: 2
v = [0, 1] 
#size: 2X3
M = [[0, 1, 2],
     [3, 4, 5]]
#size: 2X2X3
v*M = [[[0, 0, 0],
        [0, 0, 0]],
       [[0, 1, 2],
        [3, 4, 5]]]

For two vectors v1 and v2, I can use torch.bmm(v1.view(1, -1, 1), v2.view(1, 1, -1)). This can be easily extended for a batch of vectors. However, I am not able to find a solution for vector-matrix case. Also, I need to do this operation for batches of vectors and matrices.

Answer 1

You can use torch.einsum operator:

torch.einsum('bp,bqr->bpqr', v, M) # batch-wise operation v.shape=(b,p) M.shape=(b,q,r)
torch.einsum('p,qr->pqr', v, M)    # cross-batch operation