I have four multidimensional tensors v[i,j,k], a[i,s,l], w[j,s,t,m], x[k,t,n] in Numpy, and I am trying to compute the tensor z[l,m,n] given by:
z[l,m,n] = sum_{i,j,k,s,t} v[i,j,k] * a[i,s,l] * w[j,s,t,m] * x[k,t,n]
All the tensors are relatively small (say less that 32k elements in total), however I need to perform this computation many times, so I would like the function to have as little overhead as possible.
I tried to implement it using numpy.einsum like this:
z = np.einsum('ijk,isl,jstm,ktn', v, a, w, x)
but it was very slow. I also tried the following sequence of numpy.tensordot calls:
z = np.zeros((a.shape[-1],w.shape[-1],x.shape[-1]))
for s in range(a.shape[1]):
for t in range(x.shape[1]):
res = np.tensordot(v, a[:,s,:], (0,0))
res = np.tensordot(res, w[:,s,t,:], (0,0))
z += np.tensordot(res, x[:,s,:], (0,0))
inside of a double for loop to sum over s and t (both s and t are very small, so that is not too much of a problem). This worked much better, but it is still not as fast as I would expect. I think this may be because of all the operations that tensordot needs to perform internally before taking the actual product (e.g. permuting the axes).
I was wondering if there is a more efficient way to implement this kind of operations in Numpy. I also wouldn't mind implementing this part in Cython, but I'm not sure what would be the right algorithm to use.