2

Let's say I have two 3D arrays A and Bof shape (3, 4, N) and (4, 3, N).

I can compute the dot product between slices along the third axis with

with_einsum = np.eisum('ikl,kjl->ijl', A, B)

Is it possible to perform the same operation with numpy.tensordot ?

floflo29
  • 2,261
  • 2
  • 22
  • 45

1 Answers1

3

With np.einsum('ikl,kjl->ijl', A, B), there is axis alignment requirement with string - l that stays with the inputs and the output. As such, using np.tensordot might not necessarily result in performance improvement, but since the question has specifically asked for it, let's suggest it anyway. Now, np.tensordot would spread out the axes that don't take part in sum-reduction as separate axes, resulting in (N,N). So, to get to the final output, we need to extract the elements along the diagonal of the spread-out axes.

Here's how a solution with np.tensordot would look like -

mask = np.eye(N,dtype=bool)
out = np.tensordot(A,B,axes=((1),(0))).swapaxes(1,2)[:,:,mask]

There would be cases where np.dot/np.tensordot might come out as winner, but that requires that the sum-reduction axes have decent lengths. See this post for a more detailed analysis.

For the given problem, that's not the case as the length of sum-reduction is just 4. So, I would think einsum would be best here!

Community
  • 1
  • 1
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • So, in case of an axis which does not take part in the operation (i.e. we just perform the same operation along slices along this axis), `np.einsum` is the most suited function? – floflo29 Jan 25 '17 at 10:11
  • @floflo29 Linked another post which should shed more info on it. Check that out too! – Divakar Jan 25 '17 at 10:22
  • Is there a generic way to algebraically write the expression computed with `np.einsum`? I will also give a look at your link! – floflo29 Jan 25 '17 at 10:37
  • @floflo29 Not sure what you mean by writing algebraically. Could you elaborate? I would think that the string notation used by `einsum` is pretty close how one would implement in a loopy way and as such is pretty intuitive, once we get to know what each string means. – Divakar Jan 25 '17 at 10:40
  • For example, if we have `C=np.einsum('ikl,kjl->ijl', A, B)` we can write it: C(i,j,l) = sum(k)A(i,k,l)*B(k,j,l) – floflo29 Jan 25 '17 at 10:44
  • @floflo29 I see. Sadly, there isn't any, sorry. – Divakar Jan 25 '17 at 10:45