multiplying a vector (1 x N) by a tensor (N x M x M)

user_14013776 0
Viewed 378

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.

pytorch matrix-multiplication numpy arrays python
2 Answers

You can use simple einsum:

#this gives you 2-D array (M,M)
np.einsum('i,ijk->jk',a,b)

output:

[[38 44]
 [50 56]]

or another solution:

#this gives you 3-D array (1,M,M)
a[None,:]@b.swapaxes(0,1)

output:

[[[38 44]
  [50 56]]]
user_49759810

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->jk', A, B)
np.einsum('i,ijk->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

user_46144040
Related