Leaving dimension unaffected with numpy tensordot

Question

I've been using np.tensordot in the past without any problems, but my current example, I struggle to understand the result.

For np.tensordot(d * y, r, axes=((1, 2, 3), (2, 3, 4))).shape, I would expect a shape of (6, 5), but instead, I get (6, 6, 5). When I would run tensordot 6 times on axis0, I would however get the expected result, but I'd rather have tensordot do this for me in one call. What's wrong with this?

>>> import numpy as np
>>> d = np.random.rand(6, 7, 1, 2)
>>> y = np.random.rand(6, 7, 1, 2)
>>> r = np.random.rand(6, 5, 7, 1, 2) > 0.5
>>> 
>>> np.tensordot(d * y, r, axes=((1, 2, 3), (2, 3, 4))).shape
(6, 6, 5)
>>> np.tensordot((d * y)[0], r[0], axes=((0, 1, 2), (1, 2, 3))).shape
(5,)
>>> np.tensordot((d * y)[1], r[1], axes=((0, 1, 2), (1, 2, 3))).shape
(5,)
>>> np.tensordot((d * y)[2], r[2], axes=((0, 1, 2), (1, 2, 3))).shape
(5,)
...
>>> np.tensordot((d * y)[5], r[5], axes=((0, 1, 2), (1, 2, 3))).shape
(5,)

Answer 1

Consider a simpler case:

In [709]: d=np.ones((6,2));
In [710]: np.tensordot(d,d,axes=(1,1)).shape
Out[710]: (6, 6)

This is equivalent to:

In [712]: np.einsum('ij,kj->ik',d,d).shape
Out[712]: (6, 6)

This isn't ij,ij->i. It's a outer product on the unlisted axes, not an element by element one.

You have (6, 7, 1, 2) and (6, 5, 7, 1, 2), and want to sum on (7,1,2). It's doing an outer product on the (6) and (6,5).

np.einsum('i...,ij...->ij',d,r) would do, I think, what you want.

Under the covers, tensordot reshapes and swaps axes so that the problem becomes a 2d np.dot call. Then it reshapes and swaps back as needed.

---

Correction; I can't use ellipsis for the 'dotted' dimensions

In [726]: np.einsum('aijk,abijk->ab',d,r).shape
Out[726]: (6, 5)

and method:

In [729]: (d[:,None,...]*r).sum(axis=(2,3,4)).shape
Out[729]: (6, 5)

timings

In [734]: timeit [np.tensordot(d[i],r[i], axes=((0, 1, 2), (1, 2, 3))) for i in 
     ...: range(6)]
145 µs ± 514 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [735]: timeit np.einsum('aijk,abijk->ab',d,r)
7.22 µs ± 34.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [736]: timeit (d[:,None,...]*r).sum(axis=(2,3,4))
16.6 µs ± 84.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Another solution, using the @ (matmul) operator

In [747]: timeit np.squeeze(d.reshape(6,1,14)@r.reshape(6,5,14).transpose(0,2,1))
11.4 µs ± 28.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Answer 2

In tensordot, every axis is either:

• summed over (contracted) and thus eliminated from the final result, or
• unaffected (and unconstrained) and appears exactly once in the summation.

So when you write tensordot(d * y, r, axes=((1, 2, 3), (2, 3, 4))), you are calculating:

T[i j k] = ∑[l m n] dy[i l m n] r[j k l m n]

where dy ≡ d * y. What you want to calculate is

T[i k] = ∑[l m n] dy[i l m n] r[i k l m n]

Notice that i appears twice, yet isn't being summed over. This means i is actually constrained here (think of it as an implicit Kronecker delta). Hence, this isn't something that tensordot can do all by itself.

The simplest way is to use einsum and just explicitly declare what you want:

np.einsum("i l m n, i k l m n -> i k", d * y, r)

Since einsum can see the entire expression you're trying to calculate, it should be able to find a relatively optimal way to perform this calculation.