CODEWITHSUNDEEP
pythonnumpymatrixlinear-algebratheano
erobertc
erobertc

Reputation: 644

Product of matrix and 3-way tensor in Numpy/Theano

I would like to compute the product of a matrix X with shape (a, b) and a tensor Y with shape (a, b, c) such that the result Z has shape (a, c), and row i (i = 1...a) of Z is the product of row i of X and the matrix slice (i, :, :) of Y.

Is there a convenient way to do this in NumPy and Theano, ideally using built-in functions, and without using loops or computing unnecessary matrix products?

Upvotes: 0

Views: 82

Answers (1)

hpaulj
hpaulj

Reputation: 231550

With your description, writing an einsum expression is easy:

In [428]: X=np.arange(6).reshape(2,3)
In [429]: Y=np.arange(2*3*4).reshape(2,3,4)

In [431]: np.einsum('ab,abc->ac',X,Y)
Out[431]: 
array([[ 20,  23,  26,  29],
       [200, 212, 224, 236]])
In [432]: _.shape
Out[432]: (2, 4)

np.matmul or the @ operator is a little trickier, though probably as fast:

In [438]: (X[:,None,:]@Y).squeeze()
Out[438]: 
array([[ 20,  23,  26,  29],
       [200, 212, 224, 236]])

The intermediate stage will be (a,1,c) in shape, i.e. 

(a,1,b)@(a,b,c)=>(a,1,c)   # with sum on b

Upvotes: 1

Related Questions