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:
@mmohrhard this looks like an issue -- thank you for reporting.  I'm not sure what's going on at the moment.  I was expecting chunks to be: `(4,5)` and instead we have `((5,), (5,))`

For example, if we create a dask array with two equally sized chunks things look ok:

```python
In [127]: np_data = np.random.normal(10, 0.1, size=(20, 20))

In [128]: y = da.from_array(np_data, chunks=(10, 10))

In [129]: y.chunks
Out[129]: ((10, 10), (10, 10))
```

Right, two chunks of size `(10,10)`.  If instead we created one big chunk:

```python
In [130]: y = da.from_array(np_data, chunks=(20, 20))

In [131]: y.chunks
Out[131]: ((20,), (20,))
```

It looks like there is either a bug or dask internals is doing something I don't understand.  The latter is not surprising given that I am not very familiar with dask array internals.  @jrbourbeau  if you have some time would you care to comment ?
I've been digging in to this and I think there are a few issues here (although
possibly just manifestations of the same issue)

Following @mmohrhard's example

```python
In [1]: import numpy as np

In [2]: import dask.array as da

In [3]: a = np.ones(20).reshape(4, 5)

In [4]: b = da.from_array(a, chunks=(4, 5))

In [5]: np.tensordot(a, a, axes=2)
Out[5]: array(20.)

In [6]: da.tensordot(b, b, axes=2).compute()
<snip> #elided because same as original traceback
ValueError: Chunks do not add up to shape. Got chunks=((5,), (5,)), shape=(4, 5)
```

For a tensor double contraction we want the axes we're reducing over in sequence
to be the same size, which works if we specify that:

```python
In [7]: da.tensordot(b, b, axes=([0, 1], [0, 1])).compute()
Out[7]: 20.0
```

If `axes=2` and the chunk sizes aren't the same (e.g. `(5,5)`), then
`dask.array.core.unify_chunks` is called but the inferred axes are `ij, ji`
instead of `ij, ij`

From the debugger at this point:
```python
[(dask.array<array, shape=(4, 5), dtype=float64, chunksize=(4, 5), chunktype=numpy.ndarray>, [0, 1]), (dask.array<array, shape=(4, 5), dtype=float64, chunksize=(4, 5), chunktype=numpy.ndarray>, [1, 0])]
```

Note the `[0, 1]` and `[1, 0]` axes generated by `da.tensordot`. 


There are a few other different failure modes:

When chunks are symmetric but one axis is only one chunk wide:
```python
In [10]: a = np.ones(50).reshape(10, 5)

In [11]: b = da.from_array(a, chunks=(5, 5))

In [12]: b.chunks
Out[12]: ((5, 5), (5,))

In [13]: da.tensordot(b, b, axes=2).compute()
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-13-c60ee22c64a9> in <module>
----> 1 da.tensordot(b, b, axes=2).compute()

~/github.com/dask/dask/dask/array/routines.py in tensordot(lhs, rhs, axes)
    262         right_index[r] = left_index[l]
    263
--> 264     intermediate = blockwise(
    265         _tensordot,
    266         out_index,

~/github.com/dask/dask/dask/array/blockwise.py in blockwise(func, out_ind, name, token, dtype, adjust_chunks, new_axes, align_arrays, concatenate, meta, *args, **kwargs)
    144
    145     if align_arrays:
--> 146         chunkss, arrays = unify_chunks(*args)
    147     else:
    148         arginds = [(a, i) for (a, i) in toolz.partition(2, args) if i is not None]

~/github.com/dask/dask/dask/array/core.py in unify_chunks(*args, **kwargs)
  3173             )
  3174             if chunks != a.chunks and all(a.chunks):
-> 3175                 arrays.append(a.rechunk(chunks))
  3176             else:
  3177                 arrays.append(a)

~/github.com/dask/dask/dask/array/core.py in rechunk(self, chunks, threshold, block_size_limit)
  2158         from . import rechunk  # avoid circular import
  2159
-> 2160         return rechunk(self, chunks, threshold, block_size_limit)
  2161
  2162     @property

~/github.com/dask/dask/dask/array/rechunk.py in rechunk(x, chunks, threshold, block_size_limit)
    227     if isinstance(chunks, (tuple, list)):
    228         chunks = tuple(lc if lc is not None else rc for lc, rc in zip(chunks, x.chunks))
--> 229     chunks = normalize_chunks(
    230         chunks, x.shape, limit=block_size_limit, dtype=x.dtype, previous_chunks=x.chunks
    231     )

~/github.com/dask/dask/dask/array/core.py in normalize_chunks(chunks, shape, limit, dtype, previous_chunks)
  2467             for c, s in zip(map(sum, chunks), shape)
  2468         ):
-> 2469             raise ValueError(
  2470                 "Chunks do not add up to shape. "
  2471                 "Got chunks=%s, shape=%s" % (chunks, shape)

ValueError: Chunks do not add up to shape. Got chunks=((5, 5), (5, 5)), shape=(10, 5)
```

When chunks are symmetric but different numbers of chunks along an axis:
```python
In [14]: a = np.ones(200).reshape(20, 10)

In [15]: b = da.from_array(a, chunks=(5, 5))

In [16]: b.chunks
Out[16]: ((5, 5, 5, 5), (5, 5))

In [17]: da.tensordot(b, b, axes=2).compute()
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-17-c60ee22c64a9> in <module>
----> 1 da.tensordot(b, b, axes=2).compute()

~/github.com/dask/dask/dask/array/routines.py in tensordot(lhs, rhs, axes)
    262         right_index[r] = left_index[l]
    263
--> 264     intermediate = blockwise(
    265         _tensordot,
    266         out_index,

~/github.com/dask/dask/dask/array/blockwise.py in blockwise(func, out_ind, name, token, dtype, adjust_chunks, new_axes, align_arrays, concatenate, meta, *args, **kwargs)
    144
    145     if align_arrays:
--> 146         chunkss, arrays = unify_chunks(*args)
    147     else:
    148         arginds = [(a, i) for (a, i) in toolz.partition(2, args) if i is not None]

~/github.com/dask/dask/dask/array/core.py in unify_chunks(*args, **kwargs)
  3147     blockdim_dict = {a.name: a.chunks for a, ind in arginds if ind is not None}
  3148
-> 3149     chunkss = broadcast_dimensions(nameinds, blockdim_dict, consolidate=common_blockdim)
  3150     breakpoint()
  3151     max_parts = max(arg.npartitions for arg, ind in arginds if ind is not None)

~/github.com/dask/dask/dask/blockwise.py in broadcast_dimensions(argpairs, numblocks, sentinels, consolidate)
    777
    778     if consolidate:
--> 779         return toolz.valmap(consolidate, g2)
    780
    781     if g2 and not set(map(len, g2.values())) == set([1]):

~/miniconda3/envs/daskdev/lib/python3.8/site-packages/toolz/dicttoolz.py in valmap(func, d, factory)
    81     """
    82     rv = factory()
---> 83     rv.update(zip(iterkeys(d), map(func, itervalues(d))))
    84     return rv
    85

~/github.com/dask/dask/dask/array/core.py in common_blockdim(blockdims)
  3054
  3055     if len(set(map(sum, non_trivial_dims))) > 1:
-> 3056         raise ValueError("Chunks do not add up to same value", blockdims)
  3057
  3058     # We have multiple non-trivial chunks on this axis

ValueError: ('Chunks do not add up to same value', {(5, 5), (5, 5, 5, 5)})
```

In both cases, setting the `axes=([0, 1], [0, 1])` returns the correct result:
```python
In [19]: a = np.ones(50).reshape(10, 5)

In [20]: b = da.from_array(a, chunks=(5, 5))

In [21]: da.tensordot(b, b, axes=([0, 1], [0, 1])).compute()
Out[21]: 50.0

In [22]: a = np.ones(200).reshape(20, 10)

In [23]: b = da.from_array(a, chunks=(5, 5))

In [24]: da.tensordot(b, b, axes=([0, 1], [0, 1])).compute()
Out[24]: 200.0
```

I think this can be fixed if we change how the `arginds` are calculated, because
this also fails for the case when `axes=2` and `ndim>2`, in which case the `-2`
and `-1` axis from `a` should be reduced against the `0` and `1` axis of `b`,
respectively. This fails in `da.tensordot` if `axes=2` but not if those `axes`
are specified:

```python
In [39]: a = np.arange(60.).reshape(3,4,5)
    ...: b = np.arange(60.).reshape(4,5,3)

In [40]: np.tensordot(a, b, axes=2)
Out[40]:
array([[ 7410.,  7600.,  7790.],
      [18810., 19400., 19990.],
      [30210., 31200., 32190.]])

In [41]: c = da.from_array(a)

In [42]: d = da.from_array(b)

In [43]: da.tensordot(c, d, axes=2).compute()
<snip>
ValueError: Chunks do not add up to shape. Got chunks=((3,), (5,), (5,)), shape=(3, 4, 5)

In [44]: da.tensordot(c, d, axes=([1, 2], [0, 1])).compute()
Out[44]:
array([[ 7410.,  7600.,  7790.],
      [18810., 19400., 19990.],
      [30210., 31200., 32190.]])
```


yep, got it, fix incoming

Resolving patch:
diff --git a/dask/array/routines.py b/dask/array/routines.py
--- 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
--- 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)