Tensordot Explanation from Numpy documentation

Question

I don't udnerstand how tensordot works and I was reading the official documentation but I don't understand at all what is happening there.

a = np.arange(60.).reshape(3,4,5)
b = np.arange(24.).reshape(4,3,2)
c = np.tensordot(a,b, axes=([1,0],[0,1]))
c.shape
(5, 2)

Why is the shape (5, 2)? What exactly is happening?

I also read this article but the answer is confusing me.

 In [7]: A = np.random.randint(2, size=(2, 6, 5))
   ...:  B = np.random.randint(2, size=(3, 2, 4))
   ...:

In [9]: np.tensordot(A, B, axes=((0),(1))).shape
Out[9]: (6, 5, 3, 4)

A : (2, 6, 5) -> reduction of axis=0
B : (3, 2, 4) -> reduction of axis=1

Output : `(2, 6, 5)`, `(3, 2, 4)` ===(2 gone)==> `(6,5)` + `(3,4)` => `(6,5,3,4)`

Why is the shape (6, 5, 3, 4)?

`tensordot` is kind of dated. `einsun` and `matmul` are easier to use. But your example does the sum-of-products reduction on the first 2 dimensions, and 'outer' on the last. `np.einsum('ijk,jil->kl', A,B)` — hpaulj, Oct 31 '21 at 21:23

hpaulj · Answer 1 · 2021-10-31T23:43:57.460

In [196]: a = np.arange(60.).reshape(3,4,5)
     ...: b = np.arange(24.).reshape(4,3,2)
     ...: c = np.tensordot(a,b, axes=([1,0],[0,1]))
In [197]: c
Out[197]: 
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

I find the einsum equivalent to be easier to "read":

In [198]: np.einsum('ijk,jil->kl',a,b)
Out[198]: 
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

tensordot works by transposing and reshaping the inputs to reduce the problem to a simple dot:

In [204]: a1 = a.transpose(2,1,0).reshape(5,12)
In [205]: b1 = b.reshape(12,2)
In [206]: np.dot(a1,b1)        # or a1@b1
Out[206]: 
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

tensordot can do further manipulation to the result, but that's not needed here.

I had to try several things before I got a1/b1 right. For example a.transpose(2,0,1).reshape(5,12) produces the right shape, but different values.

yet another version:

In [210]: (a.transpose(1,0,2)[:,:,:,None]*b[:,:,None,:]).sum((0,1))
Out[210]: 
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

Tensordot Explanation from Numpy documentation

1 Answers1