Repository: dask/dask

Issue:
Decompose tensordot function to avoid memory
Consider a Matrix-Vector multiply `z = W.dot(v)` for `W` a short and fat matrix and `v` a tall vector.  Under common chunking approaches this has a single output chunk

```
     v
     v    
     v
WWWW z
```

To compute this efficiently in small space we want to compute the three chunks first `W1.dot(v1)`, `W2.dot(v2)`, `W3.dot(v3)`, allow their children to be released, and then sum these intermediate results to our final result.

Unfortunately this is not what the current `tensordot` function does.  It provides the two iterators `[W1, W2, W3]`, `[v1, v2, v3]` to a function `many`, which does both the per-block `tensordot` and the resulting sum.  Because of this choice, all of `W` and all of `v` will be in memory at once.

So we should change this behavior.  The only issue here is that the full scope of `tensordot` is rather large and it largely depends on `atop` which, while beautiful and saves us untold amounts of frustration, also isn't something that we can easily tweak.



Discussion:
@shoyer @clarkfitzg solutions to this problem may come close to your desire to support more sophisticated `map+atop` solutions.  If you have time to think about that application and what a reasonable feature set might look like, then now would be a good time to write up an issue for it (if there isn't one already.)

Looking over use cases of `atop` it may be reasonable to combine the current implementations of `atop` and `many`.  

Currently `atop` takes one function and it gives this one function blocks or iterators of blocks for each free index in the output (iterators are caused when we have dummy indices.)  

So in `x_ij + y_jk` atop calls the input function `n_i * n_k` times, each time on an iterator of length `n_j`.

One alternate proposal is to provide `atop` with two functions, a first mapping function to be called for each index possibility (both free and dummy indices), and a second function to reduce over the dummy indices.

So in `x_ij + y_jk` we would call the mapper function`n_i * n_j * n_k` times, each time on a single block per input (two inputs).  We would then call a second reduction function `n_i * n_k` times on those results, much in the same way we currently call `atop`.

On second thought.  I can achieve this now with multiple calls to `atop`.


Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1622,17 +1622,6 @@ def transpose(a, axes=None):
                 a, tuple(range(a.ndim)), dtype=a._dtype)
 
 
-@curry
-def many(a, b, binop=None, reduction=None, **kwargs):
-    """
-    Apply binary operator to pairwise to sequences, then reduce.
-
-    >>> many([1, 2, 3], [10, 20, 30], mul, sum)  # dot product
-    140
-    """
-    return reduction(map(partial(binop, **kwargs), a, b))
-
-
 alphabet = 'abcdefghijklmnopqrstuvwxyz'
 ALPHABET = alphabet.upper()
 
@@ -1658,25 +1647,29 @@ def tensordot(lhs, rhs, axes=2):
         raise NotImplementedError("Simultaneous Contractions of multiple "
                 "indices not yet supported")
 
+    if lhs._dtype is not None and rhs._dtype is not None :
+        dt = np.promote_types(lhs._dtype, rhs._dtype)
+    else:
+        dt = None
+
     left_index = list(alphabet[:lhs.ndim])
     right_index = list(ALPHABET[:rhs.ndim])
     out_index = left_index + right_index
+
     for l, r in zip(left_axes, right_axes):
         out_index.remove(right_index[r])
-        out_index.remove(left_index[l])
         right_index[r] = left_index[l]
 
-    if lhs._dtype is not None and rhs._dtype is not None :
-        dt = np.promote_types(lhs._dtype, rhs._dtype)
-    else:
-        dt = None
+    func = partial(np.tensordot, axes=(left_axes, right_axes))
+    intermediate = atop(func, out_index,
+                        lhs, left_index,
+                        rhs, right_index, dtype=dt)
+
+    int_index = list(out_index)
+    for l in left_axes:
+        out_index.remove(left_index[l])
 
-    func = many(binop=np.tensordot, reduction=sum,
-                axes=(left_axes, right_axes))
-    return atop(func,
-                out_index,
-                lhs, tuple(left_index),
-                rhs, tuple(right_index), dtype=dt)
+    return atop(sum, out_index, intermediate, int_index, dtype=dt)
 
 
 @wraps(np.dot)


Test patch: