Repository: dask/dask

Issue:
subgraph_callable failure with dask 0.20.1
This fails on python 2.7.14, dask 0.20.1,  but works on dask 0.19.4.

As far as I can tell, a task like `(operator.mul, 0j, [np.ndarray(...)])` gets constructed by `subgraph_callable` which produces the `TypeError: can't multiply sequence by non-int of type 'complex'` exception below


```python
from __future__ import print_function

import dask.array as da
import numpy as np

row_chunks = (25, 25, 25, 25)
chan_chunks = (4, 4, 4, 4)
corr1_chunks = (2,)
corr2_chunks = (2,)

rows = sum(row_chunks)
chans = sum(chan_chunks)
corr1 = sum(corr1_chunks)
corr2 = sum(corr2_chunks)

shape = (rows, chans, corr1, corr2)
chunks = (row_chunks, chan_chunks, corr1_chunks, corr2_chunks)

vis = (da.random.random(shape, chunks=chunks) +
       da.random.random(shape, chunks=chunks)*1j)


def _grid2(vis):
    vis = vis[0][0]
    return np.ones(vis.shape[:2], dtype=vis.dtype)


grid2 = da.core.atop(_grid2, ("row", "chan"),
                     vis, ("row", "chan", "corr1", "corr2"),
                     dtype=vis.dtype)


print(grid2.compute().shape)
```

Produces the following:

```traceback$ python dask_atop_issue.py 
Traceback (most recent call last):
  File "dask_atop_issue.py", line 38, in <module>
    print(grid.compute().shape)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 156, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 397, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/threaded.py", line 76, in get
    pack_exception=pack_exception, **kwargs)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/local.py", line 501, in get_async
    raise_exception(exc, tb)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/local.py", line 272, in execute_task
    result = _execute_task(task, data)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/local.py", line 253, in _execute_task
    return func(*args2)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/optimization.py", line 942, in __call__
    dict(zip(self.inkeys, args)))
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 132, in _get_recursive
    res = cache[x] = _get_recursive(dsk, dsk[x], cache)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 136, in _get_recursive
    args2 = [_get_recursive(dsk, k, cache) for k in args]
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 126, in _get_recursive
    return [_get_recursive(dsk, k, cache) for k in x]
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 132, in _get_recursive
    res = cache[x] = _get_recursive(dsk, dsk[x], cache)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 136, in _get_recursive
    args2 = [_get_recursive(dsk, k, cache) for k in args]
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 132, in _get_recursive
    res = cache[x] = _get_recursive(dsk, dsk[x], cache)
  File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 137, in _get_recursive
    return func(*args2)
TypeError: can't multiply sequence by non-int of type 'complex'

```


Discussion:
It looks to me like you want to specify concatenate=True to the atop
function.  I'm not sure why this would have changed in recent versions
though.  Is this something that you would be willing to investigate?

On Mon, Nov 12, 2018 at 4:36 AM Simon Perkins <notifications@github.com>
wrote:

> This fails on python 2.7.14, dask 0.20.1, but works on dask 0.19.4.
>
> As far as I can tell, a task like (operator.mul, 0j, [np.ndarray(...)])
> gets constructed by subgraph_callable which produces the TypeError: can't
> multiply sequence by non-int of type 'complex' exception below
>
> from __future__ import print_function
> import dask.array as daimport numpy as np
>
> row_chunks = (25, 25, 25, 25)
> chan_chunks = (4, 4, 4, 4)
> corr1_chunks = (2,)
> corr2_chunks = (2,)
>
> rows = sum(row_chunks)
> chans = sum(chan_chunks)
> corr1 = sum(corr1_chunks)
> corr2 = sum(corr2_chunks)
>
> shape = (rows, chans, corr1, corr2)
> chunks = (row_chunks, chan_chunks, corr1_chunks, corr2_chunks)
>
> vis = (da.random.random(shape, chunks=chunks) +
>        da.random.random(shape, chunks=chunks)*1j)
>
> def _grid(vis, ny=None, nx=None):
>     print("_grid")
>     return np.ones(shape=(1, ny, nx), dtype=vis[0][0][0].dtype)
>
>
> ny, nx = 1024, 1024
>
>
> grid = da.core.atop(_grid, ("row", "ny", "nx"),
>                     vis, ("row", "chan", "corr1", "corr2"),
>                     new_axes={"ny": ny, "nx": nx},
>                     adjust_chunks={"row": 1},
>                     ny=ny, nx=nx,
>                     dtype=vis.dtype)
> print(grid.compute().shape)
>
> Produces the following:
>
> Traceback (most recent call last):
>   File "dask_atop_issue.py", line 38, in <module>
>     print(grid.compute().shape)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 156, in compute
>     (result,) = compute(self, traverse=False, **kwargs)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/base.py", line 397, in compute
>     results = schedule(dsk, keys, **kwargs)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/threaded.py", line 76, in get
>     pack_exception=pack_exception, **kwargs)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/local.py", line 501, in get_async
>     raise_exception(exc, tb)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/local.py", line 272, in execute_task
>     result = _execute_task(task, data)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/local.py", line 253, in _execute_task
>     return func(*args2)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/optimization.py", line 942, in __call__
>     dict(zip(self.inkeys, args)))
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 132, in _get_recursive
>     res = cache[x] = _get_recursive(dsk, dsk[x], cache)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 136, in _get_recursive
>     args2 = [_get_recursive(dsk, k, cache) for k in args]
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 126, in _get_recursive
>     return [_get_recursive(dsk, k, cache) for k in x]
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 132, in _get_recursive
>     res = cache[x] = _get_recursive(dsk, dsk[x], cache)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 136, in _get_recursive
>     args2 = [_get_recursive(dsk, k, cache) for k in args]
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 132, in _get_recursive
>     res = cache[x] = _get_recursive(dsk, dsk[x], cache)
>   File "/home/sperkins/tmp/python/dask/venv/local/lib/python2.7/site-packages/dask/core.py", line 137, in _get_recursive
>     return func(*args2)
> TypeError: can't multiply sequence by non-int of type 'complex'
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4201>, or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszFXA4-9q6IZ2cAUU5IvCj11LpgV5ks5uuUEugaJpZM4YZRUs>
> .
>

@mrocklin I've reduced the example further. `concatenate=True` does indeed ameliorate this issue, but its preferable not to rely on this for performance reasons. Unfortunately I don't have time to delve deep into the new dask internals.
Here is a perhaps even simpler failure that shows the problem you're running into

```python
import dask.array as da
import numpy as np

x = da.ones((10, 10), chunks=(5, 5))

def f(chunk):
    assert isinstance(chunk, np.ndarray)


y = da.core.atop(f, 'i',
                 x, 'ij',
                 dtype=x.dtype)
y.compute()
```

```python-traceback
      6 def f(chunk):
----> 7     assert isinstance(chunk, np.ndarray)
      8 
      9 

AssertionError: 
```

You are missing dimensions in your output, so they get concatenated into a list of elements.  If you want to treat that list as a numpy array then you should pass `concatenate=True`.  Really the question here is why this would have worked before.

I'm curious, how is passing `concatenate=True` a performance issue for you?
> You are missing dimensions in your output, so they get concatenated into a list of elements. If you want to treat that list as a numpy array then you should pass concatenate=True

Correct this is a mistake in my updated example, I've re-updated it, and the complex multiplied by list issue still occurs.

> I'm curious, how is passing concatenate=True a performance issue for you?

I use this style based on @shoyer's comments in https://github.com/dask/dask/pull/3412#issuecomment-381742007
Hrm, interesting, I've reduced this to the following minimal example

```python
import dask.array as da
import numpy as np

x = da.ones((10, 10), chunks=(5, 5)) * 1j


def f(L):
    pass  # never actually get here

y = da.core.atop(f, 'i',
                 x, 'ij',
                 dtype=x.dtype)
                 
y.compute()
```

Yeah, looks like a genuine bug.  My guess is that this comes up in the atop fusion code, my guess is somewhere around here: 

https://github.com/dask/dask/blob/df4c6dc05461c77d568e747dfceb5043c065fb4d/dask/array/top.py#L599-L617

If anyone has time to take a look that would be welcome.  I'm busy the next couple days.

Resolving patch:
diff --git a/dask/array/top.py b/dask/array/top.py
--- a/dask/array/top.py
+++ b/dask/array/top.py
@@ -601,20 +601,32 @@ def optimize_atop(full_graph, keys=()):
             deps = set(top_layers)
             while deps:  # we gather as many sub-layers as we can
                 dep = deps.pop()
-                if (dep in layers and
-                        isinstance(layers[dep], TOP) and
-                        not (dep != layer and dep in keep) and  # output layer
-                        layers[dep].concatenate == layers[layer].concatenate):  # punt on mixed concatenate
-                    top_layers.add(dep)
-
-                    # traverse further to this child's children
-                    for d in full_graph.dependencies.get(dep, ()):
-                        if len(dependents[d]) <= 1:
-                            deps.add(d)
-                        else:
-                            stack.append(d)
-                else:
+                if dep not in layers:
+                    stack.append(dep)
+                    continue
+                if not isinstance(layers[dep], TOP):
+                    stack.append(dep)
+                    continue
+                if (dep != layer and dep in keep):
+                    stack.append(dep)
+                    continue
+                if layers[dep].concatenate != layers[layer].concatenate:
                     stack.append(dep)
+                    continue
+
+                # passed everything, proceed
+                top_layers.add(dep)
+
+                # traverse further to this child's children
+                for d in full_graph.dependencies.get(dep, ()):
+                    # Don't allow reductions to proceed
+                    output_indices = set(layers[dep].output_indices)
+                    input_indices = {i for _, ind in layers[dep].indices if ind for i in ind}
+
+                    if len(dependents[d]) <= 1 and output_indices.issuperset(input_indices):
+                        deps.add(d)
+                    else:
+                        stack.append(d)
 
             # Merge these TOP layers into one
             new_layer = rewrite_atop([layers[l] for l in top_layers])


Test patch:
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -269,7 +269,7 @@ def f(x):
     assert z.chunks == ((7,), (7,), (2, 2, 1))
     assert_eq(z, np.ones((7, 7, 5)))
 
-    w = atop(lambda x: x[:, 0, 0], 'a', z, 'abc', dtype=x.dtype, concatenate=concatenate)
+    w = atop(lambda x: x[:, 0, 0], 'a', z, 'abc', dtype=x.dtype, concatenate=True)
     assert w.chunks == ((7,),)
     assert_eq(w, np.ones((7,)))
 
@@ -413,3 +413,16 @@ def foo(A):
 
     D = da.Array(array_dsk, name, chunks, dtype=A.dtype)
     D.sum(axis=0).compute()
+
+
+def test_dont_merge_before_reductions():
+    x = da.ones(10, chunks=(5,))
+    y = da.atop(inc, 'i', x, 'i', dtype=x.dtype)
+    z = da.atop(sum, '', y, 'i', dtype=y.dtype)
+    w = da.atop(sum, '', z, '', dtype=y.dtype)
+
+    dsk = optimize_atop(w.dask)
+
+    assert len([d for d in dsk.dicts.values() if isinstance(d, TOP)]) == 2
+
+    z.compute()