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.

Ivan · Accepted Answer · 2022-02-17 17:04:08Z

13

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

edited Feb 17, 2022 at 17:04

Ivan

41.3k9 gold badges78 silver badges120 bronze badges

answered Mar 8, 2019 at 19:02

Separius

1,3069 silver badges25 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Chandrahas Over a year ago

Thanks!! This also seems to faster than the matrix reshaping approach.

Chandrahas · Accepted Answer · 2019-03-08 09:42:31Z

I was able to do it with following code.

Single vector and matrix

v = torch.arange(3) M = torch.arange(8).view(2, 4) # v: tensor([0, 1, 2]) # M: tensor([[0, 1, 2, 3], # [4, 5, 6, 7]]) torch.mm(v.unsqueeze(1), M.view(1, 2*4)).view(3,2,4) tensor([[[ 0, 0, 0, 0], [ 0, 0, 0, 0]], [[ 0, 1, 2, 3], [ 4, 5, 6, 7]], [[ 0, 2, 4, 6], [ 8, 10, 12, 14]]])

For a batch of vectors and matrices, it can be easily extended using torch.bmm.

v = torch.arange(batch_size*2).view(batch_size, 2) M = torch.arange(batch_size*3*4).view(batch_size, 3, 4) torch.bmm(v.unsqueeze(2), M.view(-1, 1, 3*4)).view(-1, 2, 3, 4)

ntipakos · Accepted Answer · 2019-03-08 10:08:04Z

If [batch_size, z, x, y] is the shape of the target matrix, another solution is building two matrices of this shape with appropriate elements in each position and then apply an elementwise multiplication. It works fine with batch of vectors:

# input matrices batch_size = 2 x1 = torch.Tensor([0,1]) x2 = torch.Tensor([[0,1,2], [3,4,5]]) x1 = x1.unsqueeze(0).repeat((batch_size, 1)) x2 = x2.unsqueeze(0).repeat((batch_size, 1, 1)) # dimensions b = x1.shape[0] z = x1.shape[1] x = x2.shape[1] y = x2.shape[2] # solution mat1 = x1.reshape(b, z, 1, 1).repeat(1, 1, x, y) mat2 = x2.reshape(b,1,x,y).repeat(1, z, 1, 1) mat1*mat2

Collectives™ on Stack Overflow

Pytorch batch matrix vector outer product

3 Answers 3

1 Comment

Comments

Comments

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

Comments

Comments

Linked

Related