Repository: dask/dask

Issue:
Dask interprets boolean array as an index array
**What happened**: setting items on a dask array with a bool array incorrectly casts the array to an index array.

**What you expected to happen**: no broadcastig error.

**Minimal Complete Verifiable Example**:
```python
>>> import dask
>>> import dask.array as da
>>> arr_10_20 = da.random.random((10, 20))
>>> mask_10 = da.from_array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1]).astype(bool)
>>> val_5_20 = da.random.random((5, 20))
>>> arr_10_20[mask_10, :] = val_5_20
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
...
ValueError: Can't broadcast data with shape (5, 20) across shape (10, 20)
```


**Anything else we need to know?**:

**Environment**:

- Dask version: 2021.12.0
- Python version: 3.9
- Operating System: linux
- Install method (conda, pip, source): pip


Discussion:
Thanks for writing this up! I am able to reproduce, but I'm having trouble tracking down the exact issue. I think it is likely in this funciton: https://github.com/dask/dask/blob/e4799c0498b5e5877705bb5542d8d01116ee1320/dask/array/slicing.py#L1350 Maybe @davidhassell will be able to get to the bottom of it :)

As a workaround, using a computed array as the key for the set item seems to work as expected:

```python
arr_10_20[mask_10.compute(), :] = val_5_20
```

This may be a duplicate of #8460
Thanks for mentioning that @pavithraes they are indeed very similar, but follow different code paths. Someone might be able to fix them both at once though!
I'm happy to have to a look at this.

Patch:
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -1375,24 +1375,32 @@ def parse_assignment_indices(indices, shape):
     parsed_indices : `list`
         The reformated indices that are equivalent to the input
         indices.
-    indices_shape : `list`
-        The shape implied by of the parsed indices. For instance,
-        indices of ``(slice(0,2), 5, [4,1,-1])`` will have shape
+    implied_shape : `list`
+        The shape implied by the parsed indices. For instance, indices
+        of ``(slice(0,2), 5, [4,1,-1])`` will have implied shape
         ``[2,3]``.
     reverse : `list`
         The positions of the dimensions whose indices in the
         parsed_indices output are reversed slices.
+    implied_shape_positions: `list`
+        The positions of the dimensions whose indices contribute to
+        the implied_shape. For instance, indices of ``(slice(0,2), 5,
+        [4,1,-1])`` will have implied_shape ``[2,3]`` and
+        implied_shape_positions ``[0,2]``.
 
     Examples
     --------
     >>> parse_assignment_indices((slice(1, -1),), (8,))
-    ([slice(1, 7, 1)], [6], [])
+    ([slice(1, 7, 1)], [6], [], [0])
 
     >>> parse_assignment_indices(([1, 2, 6, 5],), (8,))
-    ([array([1, 2, 6, 5])], [4], [])
+    ([array([1, 2, 6, 5])], [4], [], [0])
 
     >>> parse_assignment_indices((3, slice(-1, 2, -1)), (7, 8))
-    ([3, slice(3, 8, 1)], [5], [1])
+    ([3, slice(3, 8, 1)], [5], [1], [1])
+
+    >>> parse_assignment_indices((slice(-1, 2, -1), 3, [1, 2]), (7, 8, 9))
+    ([slice(3, 7, 1), 3, array([1, 2])], [4, 2], [0], [0, 2])
 
     """
     if not isinstance(indices, tuple):
@@ -1417,7 +1425,8 @@ def parse_assignment_indices(indices, shape):
             )
 
     # Inititalize output variables
-    indices_shape = []
+    implied_shape = []
+    implied_shape_positions = []
     reverse = []
     parsed_indices = list(normalize_index(indices, shape))
 
@@ -1462,12 +1471,13 @@ def parse_assignment_indices(indices, shape):
             start, stop, step = index.indices(size)
             div, mod = divmod(stop - start, step)
             if div <= 0:
-                indices_shape.append(0)
+                implied_shape.append(0)
             else:
                 if mod != 0:
                     div += 1
 
-                indices_shape.append(div)
+                implied_shape.append(div)
+                implied_shape_positions.append(i)
 
         elif isinstance(index, (int, np.integer)):
             # Index is an integer
@@ -1504,14 +1514,18 @@ def parse_assignment_indices(indices, shape):
 
             # Posify an integer dask array (integer numpy arrays were
             # posified in `normalize_index`)
-            if is_dask_collection(index) and index.dtype != bool:
-                index = np.where(index < 0, index + size, index)
+            if is_dask_collection(index):
+                if index.dtype == bool:
+                    index_size = np.nan
+                else:
+                    index = np.where(index < 0, index + size, index)
 
-            indices_shape.append(index_size)
+            implied_shape.append(index_size)
+            implied_shape_positions.append(i)
 
         parsed_indices[i] = index
 
-    return parsed_indices, indices_shape, reverse
+    return parsed_indices, implied_shape, reverse, implied_shape_positions
 
 
 def concatenate_array_chunks(x):
@@ -1788,14 +1802,16 @@ def value_indices_from_1d_int_index(dim, vsize, loc0, loc1):
     value_ndim = len(value_shape)
 
     # Reformat input indices
-    indices, indices_shape, reverse = parse_assignment_indices(indices, array_shape)
+    indices, implied_shape, reverse, implied_shape_positions = parse_assignment_indices(
+        indices, array_shape
+    )
 
     # Empty slices can only be assigned size 1 values
-    if 0 in indices_shape and value_shape and max(value_shape) > 1:
+    if 0 in implied_shape and value_shape and max(value_shape) > 1:
         raise ValueError(
             f"shape mismatch: value array of shape {value_shape} "
             "could not be broadcast to indexing result "
-            f"of shape {tuple(indices_shape)}"
+            f"of shape {tuple(implied_shape)}"
         )
 
     # Set variables needed when creating the part of the assignment
@@ -1846,18 +1862,18 @@ def value_indices_from_1d_int_index(dim, vsize, loc0, loc1):
     #
     # Note that array_common_shape and value_common_shape may be
     # different if there are any size 1 dimensions being brodacast.
-    offset = len(indices_shape) - value_ndim
+    offset = len(implied_shape) - value_ndim
     if offset >= 0:
         # The array has the same number or more dimensions than the
         # assignment value
-        array_common_shape = indices_shape[offset:]
+        array_common_shape = implied_shape[offset:]
         value_common_shape = value_shape
         value_offset = 0
         reverse = [i - offset for i in reverse if i >= offset]
     else:
         # The assigmment value has more dimensions than the array
         value_offset = -offset
-        array_common_shape = indices_shape
+        array_common_shape = implied_shape
         value_common_shape = value_shape[value_offset:]
         offset = 0
 
@@ -1865,12 +1881,29 @@ def value_indices_from_1d_int_index(dim, vsize, loc0, loc1):
         if value_shape[:value_offset] != (1,) * value_offset:
             raise ValueError(
                 "could not broadcast input array from shape"
-                f"{value_shape} into shape {tuple(indices_shape)}"
+                f"{value_shape} into shape {tuple(implied_shape)}"
             )
 
     base_value_indices = []
     non_broadcast_dimensions = []
-    for i, (a, b) in enumerate(zip(array_common_shape, value_common_shape)):
+
+    for i, (a, b, j) in enumerate(
+        zip(array_common_shape, value_common_shape, implied_shape_positions)
+    ):
+        index = indices[j]
+        if is_dask_collection(index) and index.dtype == bool:
+            if math.isnan(b) or b <= index.size:
+                base_value_indices.append(None)
+                non_broadcast_dimensions.append(i)
+            else:
+                raise ValueError(
+                    f"shape mismatch: value array dimension size of {b} is "
+                    "greater then corresponding boolean index size of "
+                    f"{index.size}"
+                )
+
+            continue
+
         if b == 1:
             base_value_indices.append(slice(None))
         elif a == b:
@@ -1880,11 +1913,12 @@ def value_indices_from_1d_int_index(dim, vsize, loc0, loc1):
             base_value_indices.append(None)
             non_broadcast_dimensions.append(i)
         else:
-            # Can't check  ...
             raise ValueError(
-                f"Can't broadcast data with shape {value_common_shape} "
-                f"across shape {tuple(indices_shape)}"
+                f"shape mismatch: value array of shape {value_shape} "
+                "could not be broadcast to indexing result of shape "
+                f"{tuple(implied_shape)}"
             )
+
     # Translate chunks tuple to a set of array locations in product
     # order
     chunks = array.chunks
@@ -1985,7 +2019,6 @@ def value_indices_from_1d_int_index(dim, vsize, loc0, loc1):
                 # Index is a 1-d array
                 is_bool = index.dtype == bool
                 block_index = block_index_from_1d_index(dim, loc0, loc1, is_bool)
-
                 if is_bool:
                     block_index_size = block_index_shape_from_1d_bool_index(
                         dim, loc0, loc1


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
@@ -3855,6 +3855,15 @@ def test_setitem_extended_API_1d():
             ),
             -1,
         ],
+        [
+            (
+                slice(2, 4),
+                da.from_array(
+                    [False, False, True, True, False, False, True, False, False, True]
+                ),
+            ),
+            [[-100, -101, -102, -103], [-200, -201, -202, -203]],
+        ],
     ],
 )
 def test_setitem_extended_API_2d(index, value):
@@ -4024,17 +4033,23 @@ def test_setitem_errs():
     with pytest.raises(ValueError):
         x[[True, True, True, False], 0] = [2, 3]
 
-    x = da.ones((4, 4), chunks=(2, 2))
     with pytest.raises(ValueError):
-        x[0, da.from_array([True, False, False, True])] = [2, 3, 4]
+        x[0, [True, True, True, False]] = [2, 3]
 
-    x = da.ones((4, 4), chunks=(2, 2))
     with pytest.raises(ValueError):
-        x[0, da.from_array([True, True, False, False])] = [2, 3, 4]
+        x[0, [True, True, True, False]] = [1, 2, 3, 4, 5]
+
+    with pytest.raises(ValueError):
+        x[da.from_array([True, True, True, False]), 0] = [1, 2, 3, 4, 5]
+
+    with pytest.raises(ValueError):
+        x[0, da.from_array([True, False, False, True])] = [1, 2, 3, 4, 5]
+
+    with pytest.raises(ValueError):
+        x[:, 0] = [2, 3, 4]
 
-    x = da.ones((4, 4), chunks=(2, 2))
     with pytest.raises(ValueError):
-        x[da.from_array([True, True, True, False]), 0] = [2, 3]
+        x[0, :] = [1, 2, 3, 4, 5]
 
     x = da.ones((4, 4), chunks=(2, 2))
 
diff --git a/dask/array/tests/test_cupy_core.py b/dask/array/tests/test_cupy_core.py
--- a/dask/array/tests/test_cupy_core.py
+++ b/dask/array/tests/test_cupy_core.py
@@ -537,17 +537,23 @@ def test_setitem_errs():
     with pytest.raises(ValueError):
         x[[True, True, True, False], 0] = [2, 3]
 
-    x = da.ones((4, 4), chunks=(2, 2))
     with pytest.raises(ValueError):
-        x[0, da.from_array([True, False, False, True])] = [2, 3, 4]
+        x[0, [True, True, True, False]] = [2, 3]
 
-    x = da.ones((4, 4), chunks=(2, 2))
     with pytest.raises(ValueError):
-        x[0, da.from_array([True, True, False, False])] = [2, 3, 4]
+        x[0, [True, True, True, False]] = [1, 2, 3, 4, 5]
+
+    with pytest.raises(ValueError):
+        x[da.from_array([True, True, True, False]), 0] = [1, 2, 3, 4, 5]
+
+    with pytest.raises(ValueError):
+        x[0, da.from_array([True, False, False, True])] = [1, 2, 3, 4, 5]
+
+    with pytest.raises(ValueError):
+        x[:, 0] = [2, 3, 4]
 
-    x = da.ones((4, 4), chunks=(2, 2))
     with pytest.raises(ValueError):
-        x[da.from_array([True, True, True, False]), 0] = [2, 3]
+        x[0, :] = [1, 2, 3, 4, 5]
 
     x = da.ones((4, 4), chunks=(2, 2))