Theano: How to take a "matrix outer product" where the elements are matrices

Question

Basically, I have 2 tensors: A, where A.shape = (N, H, D), and B, where B.shape = (K, H, D). What I would like to do is to get a tensor, C, with shape (N, K, D, H) such that :

C[i, j, :, :] = A[i, :, :] * B[j, :, :]. 

Can this be done efficiently in Theano?

Side note: The actual end result that I would like to achieve is to have a tensor, E, of shape (N, K, D) such that :

E[i, j, :] = (A[i, :, :]*B[j, :, :]).sum(0)

So, if there is a way to get this directly, I would prefer it (saves on space hopefully).

Answer 1

One approach could be suggested that uses broadcasting -

(A[:,None]*B).sum(2)

Please note that the intermediate array being created would be of shape (N, K, H, D) before sum-reduction on axis=2 reduces it to (N,K,D).