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:

Patch:
diff --git a/dask/array/top.py b/dask/array/top.py
index bbd601f5c..f71f588ae 100644
--- 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])
diff --git a/docs/source/dataframe-create.rst b/docs/source/dataframe-create.rst
index c8476a03b..118990454 100644
--- a/docs/source/dataframe-create.rst
+++ b/docs/source/dataframe-create.rst
@@ -1,19 +1,19 @@
 Create and Store Dask DataFrames
 ================================
 
-Dask can create dataframes from various data storage formats like CSV, HDF,
-Apache Parquet, and others.  For most formats this data can live on various
+Dask can create DataFrames from various data storage formats like CSV, HDF,
+Apache Parquet, and others.  For most formats, this data can live on various
 storage systems including local disk, network file systems (NFS), the Hadoop
 File System (HDFS), and Amazon's S3 (excepting HDF, which is only available on
 POSIX like file systems).
 
 See the :doc:`Overview section <dataframe-overview>` for an in depth
-discussion of ``dask.dataframe`` scope, use, limitations.
+discussion of ``dask.dataframe`` scope, use, and limitations.
 
 API
 ---
 
-The following functions provide access to convert between Dask Dataframes,
+The following functions provide access to convert between Dask DataFrames,
 file formats, and other Dask or Python collections.
 
 .. currentmodule:: dask.dataframe
@@ -51,9 +51,9 @@ Pandas:
 Locations
 ---------
 
-For text, CSV, and Apache Parquet formats data can come from local disk, from
-the Hadoop File System, from S3FS, or others, by prepending the filenames with
-a protocol.
+For text, CSV, and Apache Parquet formats, data can come from local disk,
+the Hadoop File System, S3FS, or other sources, by prepending the filenames with
+a protocol:
 
 .. code-block:: python
 
@@ -61,10 +61,10 @@ a protocol.
    >>> df = dd.read_csv('hdfs:///path/to/my-data-*.csv')
    >>> df = dd.read_csv('s3://bucket-name/my-data-*.csv')
 
-For remote systems like HDFS or S3 credentials may be an issue.  Usually these
-are handled by configuration files on disk (such as a ``.boto`` file for S3)
+For remote systems like HDFS or S3, credentials may be an issue.  Usually, these
+are handled by configuration files on disk (such as a ``.boto`` file for S3),
 but in some cases you may want to pass storage-specific options through to the
-storage backend.  You can do this with the ``storage_options=`` keyword.
+storage backend.  You can do this with the ``storage_options=`` keyword:
 
 .. code-block:: python
 
@@ -75,11 +75,11 @@ storage backend.  You can do this with the ``storage_options=`` keyword.
 Dask Delayed
 ------------
 
-For more complex situations not covered by the functions above you may want to
-use :doc:`dask.delayed<delayed>` , which lets you construct
-Dask.dataframes out of arbitrary Python function calls that load dataframes.
-This can allow you to handle new formats easily, or bake in particular logic
-around loading data if, for example, your data is stored with some special
+For more complex situations not covered by the functions above, you may want to
+use :doc:`dask.delayed<delayed>`, which lets you construct Dask DataFrames out 
+of arbitrary Python function calls that load DataFrames.  This can allow you to 
+handle new formats easily or bake in particular logic around loading data if, 
+for example, your data is stored with some special format.
 
 See :doc:`documentation on using dask.delayed with
 collections<delayed-collections>` or an `example notebook
@@ -87,36 +87,36 @@ collections<delayed-collections>` or an `example notebook
 how to create a Dask DataFrame from a nested directory structure of Feather
 files (as a stand in for any custom file format).
 
-Dask.delayed is particularly useful when simple ``map`` operations aren't
+Dask delayed is particularly useful when simple ``map`` operations aren't
 sufficient to capture the complexity of your data layout.
 
 
 From Raw Dask Graphs
 --------------------
 
-This section is mainly for developers wishing to extend dask.dataframe.  It
+This section is mainly for developers wishing to extend ``dask.dataframe``.  It
 discusses internal API not normally needed by users.  Everything below can be
 done just as effectively with :doc:`dask.delayed<delayed>`  described
-just above.  You should never need to create a dataframe object by hand.
+just above.  You should never need to create a DataFrame object by hand.
 
 To construct a DataFrame manually from a dask graph you need the following
 information:
 
-1.  dask: a dask graph with keys like ``{(name, 0): ..., (name, 1): ...}`` as
+1.  Dask: a Dask graph with keys like ``{(name, 0): ..., (name, 1): ...}`` as
     well as any other tasks on which those tasks depend.  The tasks
     corresponding to ``(name, i)`` should produce ``pandas.DataFrame`` objects
-    that correspond to the columns and divisions information discussed below.
-2.  name: The special name used above
-3.  columns: A list of column names
-4.  divisions: A list of index values that separate the different partitions.
-    Alternatively, if you don't know the divisions (this is common) you can
+    that correspond to the columns and divisions information discussed below
+2.  Name: the special name used above
+3.  Columns: a list of column names
+4.  Divisions: a list of index values that separate the different partitions.
+    Alternatively, if you don't know the divisions (this is common), you can
     provide a list of ``[None, None, None, ...]`` with as many partitions as
-    you have plus one.  For more information see the Partitions section in the
-    :doc:`dataframe documentation <dataframe>`.
+    you have plus one.  For more information, see the Partitions section in the
+    :doc:`DataFrame documentation <dataframe>`
 
 As an example, we build a DataFrame manually that reads several CSV files that
-have a datetime index separated by day.  Note, you should never do this.  The
-``dd.read_csv`` function does this for you.
+have a datetime index separated by day.  Note that you should **never** do this.  
+The ``dd.read_csv`` function does this for you:
 
 .. code-block:: Python
 
diff --git a/docs/source/dataframe-performance.rst b/docs/source/dataframe-performance.rst
index e1d8c16cd..3bde21b1a 100644
--- a/docs/source/dataframe-performance.rst
+++ b/docs/source/dataframe-performance.rst
@@ -190,7 +190,7 @@ Store Data in Apache Parquet Format
 -----------------------------------
 
 HDF5 is a popular choice for Pandas users with high performance needs.  We
-encourage Dask.dataframe users to :doc:`store and load data dataframe-create>`
+encourage Dask.dataframe users to :doc:`store and load data <dataframe-create>`
 using Parquet instead.  `Apache Parquet <http://parquet.apache.org/>`_ is a
 columnar binary format that is easy to split into multiple files (easier for
 parallel loading) and is generally much simpler to deal with than HDF5 (from


Test patch:
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
index db243f08b..3f5e75f63 100644
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -269,7 +269,7 @@ def test_atop_stacked_new_axes_front(concatenate):
     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 test_gh_4176():
 
     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()