Matrix triple product with theano

Question

This is pretty much the same question as here Matrix/Tensor Triple Product? , but for theano.

So I have three matrices A , B , C of sizes n*r , m*r , l*r , and I want to compute the 3D tensor of shape (n,m,l) resulting from the triple (trilinear) product:

X[i,j,k] = \sum_a A[i,a] B[j,a] C[k,a]

A , B and C are shared variables:

A = theano.shared(numpy.random.randn(n,r))
B = theano.shared(numpy.random.randn(m,r))
C = theano.shared(numpy.random.randn(l,r))

I'd like to write it with a single theano expression, is there a way to do so? If there are many, which one is the fastest?

Answer 1

np.einsum('nr,mr,lr->nml', A, B, C)

is equivalent to

np.dot(A[:, None, :] * B[None, :, :], C.T)

which can be implemented in Theano as

theano.dot(A[:, None, :] * B[None, :, :], C.T)