0

I am not very familiar with tensor algebra and I am having trouble understanding how to make numpy.tensordot do what I want.

The example I am working with is simple: given a tensor a with shape (2,2,3) and another b with shape (2,1,3), I want a result tensor c with shape (2,1). This tensor would be the result of the following, equivalent python code:

n = a.shape[2]
c = np.zeros((2,n))
for k in range(n):
    c += a[:,:,k]*b[:,:,k]

The documentation says that the optional parameter axes:

If an int N, sum over the last N axes of a and the first N axes of b in order. The sizes of the corresponding axes must match.

But I don't understand which "axes" are needed here (furthermore, when axes is a tuple or a tuple of tuples it gets even more confusing). Examples aren't very clear to me either.

Adrian
  • 755
  • 9
  • 17

1 Answers1

1

The way tensordot works, it won't work here (not at least directly) because of the alignment requirement along the first axes. You can use np.einsum though to solve your case -

c = np.einsum('ijk,ilk->ij',a,b)

Alternatively, use np.matmul/@-operator (Python 3.x) -

np.matmul(a,b.swapaxes(1,2))[...,0] # or (a @ b.swapaxes(1,2))[...,0]
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • Just to be sure, if the b matrix changed to `(2,3)`, then `c = np.einsum('ijk,ik->ij',a,b)` would still do the same? (in which case I could just squash axes with dim = 1) – Adrian Sep 03 '19 at 13:30
  • @Adrian Yup, should be the same. – Divakar Sep 03 '19 at 13:31