Repository: dask/dask

Issue:
Common chunks in dask.array atop
Forking from https://github.com/dask/dask/issues/417

For operations backed by `atop` we currently require that chunks either match exactly, or that some are singletons.

``` python
>>> common_blockdim([(3,), (2, 1)])
set([(2, 1)])

>>> common_blockdim([(2, 1), (1, 2)])
ValueError: Chunks do not align {(1, 2), (2, 1)}

>>> common_blockdim([(3, 3), (2, 1, 3)])  # We should probably support this regardless
ValueError: Chunks do not align {(2, 1, 3), (3, 3)}
```
-  Should we relax this?  
-  If so, what is the right way to do this?
-  Should we warn in bad situations?  This opens up opportunities for people to shoot themselves in the foot.

cc @shoyer @jcrist 



Discussion:
I am reluctant to relax these requirements by much short of a major overhaul of dask.array. There is no avoiding the need to understand chunking when using the current version of dask.array. Without knowledge of the complete computation graph, it is very difficult to pick appropriate chunk sizes.


Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -205,7 +205,6 @@ def broadcast_dimensions(argpairs, numblocks, sentinels=(1, (1,)),
 
     Parameters
     ----------
-
     argpairs: iterable
         name, ijk index pairs
     numblocks: dict
@@ -217,7 +216,6 @@ def broadcast_dimensions(argpairs, numblocks, sentinels=(1, (1,)),
 
     Examples
     --------
-
     >>> argpairs = [('x', 'ij'), ('y', 'ji')]
     >>> numblocks = {'x': (2, 3), 'y': (3, 2)}
     >>> broadcast_dimensions(argpairs, numblocks)
@@ -246,8 +244,8 @@ def broadcast_dimensions(argpairs, numblocks, sentinels=(1, (1,)),
 
     g2 = dict((k, v - set(sentinels) if len(v) > 1 else v) for k, v in g.items())
 
-    if consolidate is not None:
-        g2 = valmap(consolidate, g2)
+    if consolidate:
+        return valmap(consolidate, g2)
 
     if g2 and not set(map(len, g2.values())) == set([1]):
         raise ValueError("Shapes do not align %s" % g)
@@ -1726,19 +1724,49 @@ def common_blockdim(blockdims):
     --------
 
     >>> common_blockdim([(3,), (2, 1)])
-    set([(2, 1)])
+    (2, 1)
+    >>> common_blockdim([(1, 2), (2, 1)])
+    (1, 1, 1)
     >>> common_blockdim([(2, 2), (3, 1)])  # doctest: +SKIP
     Traceback (most recent call last):
         ...
     ValueError: Chunks do not align
     """
     non_trivial_dims = set([d for d in blockdims if len(d) > 1])
-    if len(non_trivial_dims) > 1:
-        raise ValueError('Chunks do not align %s' % non_trivial_dims)
-    elif non_trivial_dims:
-        return non_trivial_dims
-    else:
-        return blockdims
+    if len(non_trivial_dims) == 1:
+        return first(non_trivial_dims)
+    if len(non_trivial_dims) == 0:
+        return max(blockdims, key=first)
+
+    if len(set(map(sum, non_trivial_dims))) > 1:
+        raise ValueError("Chunks do not add up to same value", blockdims)
+
+    # We have multiple non-trivial chunks on this axis
+    # e.g. (5, 2) and (4, 3)
+
+    # We create a single chunk tuple with the same total length
+    # that evenly divides both, e.g. (4, 1, 2)
+
+    # To accomplish this we walk down all chunk tuples together, finding the
+    # smallest element, adding it to the output, and subtracting it from all
+    # other elements and remove the element itself.  We stop once we have
+    # burned through all of the chunk tuples.
+    # For efficiency's sake we reverse the lists so that we can pop off the end
+    rchunks = [list(ntd)[::-1] for ntd in non_trivial_dims]
+    total = sum(first(non_trivial_dims))
+    i = 0
+
+    out = []
+    while i < total:
+        m = min(c[-1] for c in rchunks)
+        out.append(m)
+        for c in rchunks:
+            c[-1] -= m
+            if c[-1] == 0:
+                c.pop()
+        i += m
+
+    return tuple(out)
 
 
 def unify_chunks(*args):
@@ -1766,6 +1794,12 @@ def unify_chunks(*args):
 
     chunkss = broadcast_dimensions(nameinds, blockdim_dict,
                                    consolidate=common_blockdim)
+    max_parts = max(arg.npartitions for arg in args[::2])
+    nparts = np.prod(list(map(len, chunkss.values())))
+
+    if nparts >= max_parts * 10:
+        warnings.warn("Increasing number of chunks by factor of %d" %
+                      (nparts / max_parts))
     arrays = [a.rechunk(tuple(chunkss[j] if a.shape[n] > 1 else 1
                               for n, j in enumerate(i)))
               for a, i in arginds]


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -8,6 +8,7 @@
 from distutils.version import LooseVersion
 from operator import add, sub, getitem
 from threading import Lock
+import warnings
 
 from toolz import merge, countby, concat
 from toolz.curried import identity
@@ -29,7 +30,8 @@
                              broadcast_to, reshape, fromfunction,
                              blockdims_from_blockshape, store, optimize,
                              from_func, normalize_chunks, broadcast_chunks,
-                             atop, from_delayed, concatenate_axes)
+                             atop, from_delayed, concatenate_axes,
+                             common_blockdim)
 from dask.array.utils import assert_eq
 
 # temporary until numpy functions migrated
@@ -2235,3 +2237,61 @@ def f(a, b, c):
     z = atop(f, 'j', x, 'ijk', y, 'ki', y, 'ij', concatenate=True,
              dtype=x._dtype)
     assert_eq(z, np.ones(10))
+
+
+def test_common_blockdim():
+    assert common_blockdim([(5,), (5,)]) == (5,)
+    assert common_blockdim([(5,), (2, 3,)]) == (2, 3)
+    assert common_blockdim([(5, 5), (2, 3, 5)]) == (2, 3, 5)
+    assert common_blockdim([(5, 5), (2, 3, 5)]) == (2, 3, 5)
+    assert common_blockdim([(5, 2, 3), (2, 3, 5)]) == (2, 3, 2, 3)
+
+    assert common_blockdim([(1, 2), (2, 1)]) == (1, 1, 1)
+    assert common_blockdim([(1, 2, 2), (2, 1, 2), (2, 2, 1)]) == (1, 1, 1, 1, 1)
+
+
+def test_uneven_chunks_that_fit_neatly():
+    x = da.arange(10, chunks=((5, 5),))
+    y = da.ones(10, chunks=((5, 2, 3),))
+
+    assert_eq(x + y, np.arange(10) + np.ones(10))
+
+    z = x + y
+    assert z.chunks == ((5, 2, 3),)
+
+
+def test_elemwise_uneven_chunks():
+    x = da.arange(10, chunks=((4, 6),))
+    y = da.ones(10, chunks=((6, 4),))
+
+    assert_eq(x + y, np.arange(10) + np.ones(10))
+
+    z = x + y
+    assert z.chunks == ((4, 2, 4),)
+
+    x = da.random.random((10, 10), chunks=((4, 6), (5, 2, 3)))
+    y = da.random.random((4, 10, 10), chunks=((2, 2), (6, 4), (2, 3, 5)))
+
+    z = x + y
+    assert_eq(x + y, x.compute() + y.compute())
+    assert z.chunks == ((2, 2), (4, 2, 4), (2, 3, 2, 3))
+
+
+def test_uneven_chunks_atop():
+    x = da.random.random((10, 10), chunks=((2, 3, 2, 3), (5, 5)))
+    y = da.random.random((10, 10), chunks=((4, 4, 2), (4, 2, 4)))
+    z = atop(np.dot, 'ik', x, 'ij', y, 'jk', dtype=x._dtype, concatenate=True)
+    assert z.chunks == (x.chunks[0], y.chunks[1])
+
+    assert_eq(z, x.compute().dot(y))
+
+
+def test_warn_bad_rechunking():
+    x = da.ones((20, 20), chunks=(20, 1))
+    y = da.ones((20, 20), chunks=(1, 20))
+
+    with warnings.catch_warnings(record=True) as record:
+        x + y
+
+    assert record
+    assert '20' in record[0].message.args[0]