Repository: dask/dask
Pull request: Unify chunks in broadcast_arrays
Fixes #5754

- [ ] Tests added / passed
- [ ] Passes `black dask` / `flake8 dask`


Problem statement:
numpy.broadcast_array on dask array raises ValueError("Chunks do not align").
This is a regression. The old behaviour seems to be broadcast_array forces dask to perform a compute() to return arrays instead.

The new current behaviour seems to be going through the `__array_function__` protocol and perform a dask broadcast without automatic rechunking.

I am not sure when this has changed; and I don't know if there is anything we can do about this on the dask side -- eliminating implicit calls to compute() will surely cause some minor pains like this. 

An implicit rechunking is probably appropriate in this context, but I worry the operation may not always be well-defined (I couldn't think of an example where it is not.)

The stack trace
```
  File "<__array_function__ internals>", line 6, in broadcast_arrays
  File "/dev/shm/local/lib/python3.7/site-packages/dask/array/core.py", line 1359, in __array_function__
    return da_func(*args, **kwargs)
  File "/dev/shm/local/lib/python3.7/site-packages/dask/array/core.py", line 4025, in broadcast_arrays
    chunks = broadcast_chunks(*(e.chunks for e in args))
  File "/dev/shm/local/lib/python3.7/site-packages/dask/array/core.py", line 779, in broadcast_chunks
    raise ValueError("Chunks do not align: %s" % str(step2))
ValueError: Chunks do not align: [(55276, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 100000, 13022), (8568298,)]
```

Hints from issue discussion:
Thanks @rainwoodman .  I'm not entirely sure I understand the context.  Do you have a small example that causes this error?  
Yes.

numpy.broadcast_array used to trigger .compute() (presumably via numpy.array) and did not dispatch this into dask's broadcast_arrays function.

With 2.9.0, it goes through dask's broadcast_arrays function:
```
import dask.array as da
import numpy
import dask
print('dask version', dask.__version__)
a1 = da.ones(30, chunks=(3,))
a2 = da.ones(30, chunks=(5,))
numpy.broadcast_arrays(a1, a2)

dask version 2.9.0

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-12-8bc18a86b07d> in <module>
      3 a1 = da.ones(30, chunks=(3,))
      4 a2 = da.ones(30, chunks=(5,))
----> 5 numpy.broadcast_arrays(a1, a2)

<__array_function__ internals> in broadcast_arrays(*args, **kwargs)

~/.conda/envs/bccp/lib/python3.7/site-packages/dask/array/core.py in __array_function__(self, func, types, args, kwargs)
   1357         if da_func is func:
   1358             return handle_nonmatching_names(func, args, kwargs)
-> 1359         return da_func(*args, **kwargs)
   1360 
   1361     @property

~/.conda/envs/bccp/lib/python3.7/site-packages/dask/array/core.py in broadcast_arrays(*args, **kwargs)
   4023 
   4024     shape = broadcast_shapes(*(e.shape for e in args))
-> 4025     chunks = broadcast_chunks(*(e.chunks for e in args))
   4026 
   4027     result = [broadcast_to(e, shape=shape, chunks=chunks) for e in args]

~/.conda/envs/bccp/lib/python3.7/site-packages/dask/array/core.py in broadcast_chunks(*chunkss)
    777             step2 = [c for c in step1 if c != (1,)]
    778         if len(set(step2)) != 1:
--> 779             raise ValueError("Chunks do not align: %s" % str(step2))
    780         result.append(step2[0])
    781     return tuple(result)

ValueError: Chunks do not align: [(3, 3, 3, 3, 3, 3, 3, 3, 3, 3), (5, 5, 5, 5, 5, 5)]
```

The behavior is a bit surprising -- numpy.broadcast_array on ndarray doesn't throw a ValueError like this.

Ah I see, thank you for the very clear example.

> An implicit rechunking is probably appropriate in this context

Yes, I agree that that sounds sensible.  I've tried to implement this in https://github.com/dask/dask/pull/5765 using the `unify_chunks` function, but I have to admit that I haven't thought through the problem very deeply.  If you have some time to investigate that function and review things I would welcome the help.

Commit message: Unify chunks in broadcast_arrays

Patch to review:
diff --git a/dask/array/core.py b/dask/array/core.py
index 75b7682de4a..9221cacb4ed 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -4025,6 +4025,11 @@ def broadcast_arrays(*args, **kwargs):
     if kwargs:
         raise TypeError("unsupported keyword argument(s) provided")
 
+    # Unify uneven chunking
+    inds = [list(range(x.ndim))[::-1] for x in args]
+    uc_args = list(concat(zip(args, inds)))
+    _, args = unify_chunks(*uc_args, warn=False)
+
     shape = broadcast_shapes(*(e.shape for e in args))
     chunks = broadcast_chunks(*(e.chunks for e in args))
 
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index bfe9b1e1f93..f5a774aaa9d 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -971,6 +971,20 @@ def test_broadcast_arrays():
         assert_eq(e_a_r, e_d_r)
 
 
+def test_broadcast_arrays_uneven_chunks():
+    x = da.ones(30, chunks=(3,))
+    y = da.ones(30, chunks=(5,))
+    z = np.broadcast_arrays(x, y)
+
+    assert_eq(z, z)
+
+    x = da.ones((1, 30), chunks=(1, 3))
+    y = da.ones(30, chunks=(5,))
+    z = np.broadcast_arrays(x, y)
+
+    assert_eq(z, z)
+
+
 @pytest.mark.parametrize(
     "u_shape, v_shape",
     [


Code review comments:
diff_hunk: @@ -4025,6 +4025,11 @@ def broadcast_arrays(*args, **kwargs):
     if kwargs:
         raise TypeError("unsupported keyword argument(s) provided")
 
+    # Unify uneven chunking
+    inds = [list(range(x.ndim))[::-1] for x in args]
path: dask/array/core.py
text: @user1:
```suggestion
    inds = [list(reversed(range(x.ndim))) for x in args]
```

diff_hunk: @@ -4025,6 +4025,11 @@ def broadcast_arrays(*args, **kwargs):
     if kwargs:
         raise TypeError("unsupported keyword argument(s) provided")
 
+    # Unify uneven chunking
+    inds = [list(range(x.ndim))[::-1] for x in args]
+    uc_args = list(concat(zip(args, inds)))
path: dask/array/core.py
text: @user1:
No need for the `list` call here, the splat below iterates automatically.

```suggestion
    uc_args = concat(zip(args, inds))
```