Repository: dask/dask

Issue:
Incorrect chunk size calculation in da.tensordot
Our original test case that showed that something is going wrong with `da.tensordot` in comparison to `np.tensordot`.

```
import numpy as np
import dask.array as da
from dask.array.utils import assert_eq


a = np.ones(20).reshape(4, 5)
b = da.from_array(a, chunks=(4, 5))

np_res = np.tensordot(a, a, axes=2)
da_res = da.tensordot(b, b, axes=2)

assert_eq(np_res, da_res)
```

Error message is:

```
Traceback (most recent call last):
  File "test2.py", line 10, in <module>
    da_res = da.tensordot(b, b, axes=2)
  File "/node/bud205/envs/dugwave_master2/lib/python3.6/site-packages/dask/array/routines.py", line 272, in tensordot
    axes=(left_axes, right_axes),
  File "/node/bud205/envs/dugwave_master2/lib/python3.6/site-packages/dask/array/blockwise.py", line 145, in blockwise
    chunkss, arrays = unify_chunks(*args)
  File "/node/bud205/envs/dugwave_master2/lib/python3.6/site-packages/dask/array/core.py", line 3061, in unify_chunks
    arrays.append(a.rechunk(chunks))
  File "/node/bud205/envs/dugwave_master2/lib/python3.6/site-packages/dask/array/core.py", line 2098, in rechunk
    return rechunk(self, chunks, threshold, block_size_limit)
  File "/node/bud205/envs/dugwave_master2/lib/python3.6/site-packages/dask/array/rechunk.py", line 226, in rechunk
    chunks, x.shape, limit=block_size_limit, dtype=x.dtype, previous_chunks=x.chunks
  File "/node/bud205/envs/dugwave_master2/lib/python3.6/site-packages/dask/array/core.py", line 2404, in normalize_chunks
    "Got chunks=%s, shape=%s" % (chunks, shape)
ValueError: Chunks do not add up to shape. Got chunks=((5,), (5,)), shape=(4, 5)
```

Discussion:

Patch:
diff --git a/dask/array/routines.py b/dask/array/routines.py
index 03ea568df..e4668d466 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -239,7 +239,7 @@ def tensordot(lhs, rhs, axes=2):
     if isinstance(axes, Iterable):
         left_axes, right_axes = axes
     else:
-        left_axes = tuple(range(lhs.ndim - 1, lhs.ndim - axes - 1, -1))
+        left_axes = tuple(range(lhs.ndim - axes, lhs.ndim))
         right_axes = tuple(range(0, axes))
 
     if isinstance(left_axes, Integral):


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index 952c082f7..f192caff6 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -304,6 +304,31 @@ def test_tensordot_2(axes):
     assert_eq(da.tensordot(y, y, axes=axes), np.tensordot(x, x, axes=axes))
 
 
+@pytest.mark.parametrize("chunks", ["auto", (4, 6), (2, 3), (4, 3), (2, 6)])
+def test_tensordot_double_contraction_neq2(chunks):
+    # Regression test for https://github.com/dask/dask/issues/5472
+    x = np.arange(24).reshape(4, 6)
+    y = da.from_array(x, chunks=chunks)
+    assert_eq(da.tensordot(y, y, axes=2), np.tensordot(x, x, axes=2))
+
+
+def test_tensordot_double_contraction_ngt2():
+    # Regression test for https://github.com/dask/dask/issues/5472
+    x = np.arange(60.0).reshape(3, 4, 5)
+    y = np.arange(60.0).reshape(4, 5, 3)
+    u = da.from_array(x)
+    v = da.from_array(y)
+
+    assert_eq(da.tensordot(u, v, axes=2), np.tensordot(x, y, axes=2))
+
+    x = np.arange(60.0).reshape(3, 4, 5)
+    y = np.arange(60.0).reshape(4, 5, 3)
+    u = da.from_array(x, chunks=3)
+    v = da.from_array(y)
+
+    assert_eq(da.tensordot(u, v, axes=2), np.tensordot(x, y, axes=2))
+
+
 def test_tensordot_more_than_26_dims():
     ndim = 27
     x = np.broadcast_to(1, [2] * ndim)