Repository: dask/dask

Issue:
fixed broadcast of array chunks containing a scalar during `compute()`
- [x] Closes #8822 
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`

I decided, somewhat reluctantly, to create this PR after thinking long and hard about Issue #8822 .  I was reluctant because this PR runs counter to PR #8847 .  But I hope @jsignell doesn't mind. :smile:  In any case, if this PR is not approved, then I hope it would have at least sparked a helpful discussion.

In principle, in `dask.array`, the content of any chunk must be an array whose number of dimensions matches that of the `dask` array possessing that chunk.

In practice, however, there appears to be room for grace, as under certain conditions, `dask` is able to work seamlessly with chunks whose contents violate this rule.

For example, consider the following chunk-function:
```python
import numpy as np


def f(x, ndim=0):
    """
	Ignore `x` and return arbitrary one-element array of dimensionality `ndim`
	
	For example,
	f(x, 0) = array(5)
	f(x, 1) = array([5])
	f(x, 2) = array([[5]])
	f(x, 3) = array([[[5]]])
	"""
    return np.array(5)[(np.newaxis,) * ndim]
```
which is applied, when `ndim = 1`, to the following `dask` array using `map_blocks()`: 
```python
import dask.array as da
from functools import partial


array = da.from_array([1] + [2, 2] + [3, 3, 3], chunks=((1, 2, 3),))
out_chunks = ((1, 1, 1),)

array.map_blocks(partial(f, ndim=1), chunks=out_chunks).compute()
# returns array([5, 5, 5])
```

The return value is as expected.  However, the following calls also return the same value:
```python
array.map_blocks(partial(f, ndim=2), chunks=out_chunks).compute()
# returns array([5, 5, 5])
array.map_blocks(partial(f, ndim=8), chunks=out_chunks).compute()
# returns array([5, 5, 5])
```
IMO, this attests to `dask`'s powerful ability to recognize what the dimensionality of each chunk should be and correct it.  But oddly, as first reported by @dcherian in #8822 , currently this behavior does not to extend to `ndim = 0`:
```python
array.map_blocks(partial(f, ndim=0), chunks=out_chunks).compute()
# currently raises a confusing error
```

This PR fixes this issue allowing the latter to return the same value as in the former cases.  It also fixes the following case, in which the chunk function returns a non-array-like type: 
```python
array.map_blocks(lambda x: 5, chunks=out_chunks).compute()
# currently raises a confusing error
```
allowing `map_blocks()` to return the same value as in the former cases.


Discussion:
Thanks for writing this up @ParticularMiner! I am going to take a look as soon as possible. I can say right away though that I like the idea of doing what the user expects rather than erroring. 
> IMO, this attests to dask's powerful ability to recognize what the dimensionality of each chunk should be and correct it.

This is OK but can also be confusing when a user is trying to figure out how to return the right thing. Maybe there should be an opt-in strict mode (default off).
> > IMO, this attests to dask's powerful ability to recognize what the dimensionality of each chunk should be and correct it.
> 
> This is OK but can also be confusing when a user is trying to figure out how to return the right thing. Maybe there should be an opt-in strict mode (default off).

@dcherian can you help me understand what "the right thing" would be in this case? I am pretty convinced that in your original example the right thing is `[1., 2., 3.]`, but maybe I misunderstood something.
I think in the strictest sense, the "right thing" is that the function should return a 1D array because that's what dask expects. So for that example I needed to specify `keepdims=True`.

The convenient thing is to add/remove dimensions as needed which is what dask does by default.

My point is that it's hard to understand how map_blocks works when the default behaviour is to do the "convenient thing". The following is totally unexpected to me =)
```
array.map_blocks(partial(f, ndim=8), chunks=out_chunks).compute()
# returns array([5, 5, 5])
```

To be clear, I think this PR makes sense given dask's  current "convenient behaviour". 

For understanding map_blocks, it'd be nice to  have an opt-in strict mode. In fact, I am surprised that `assert_eq` in the tests works though. I thought it did strict checking of shapes per block.
Ah ok that makes sense. It would probably be useful to have a strict config value for dask more generally. 
I'm going to wait to merge this until after the release on Friday. I want to make sure to give people time to react. 
> This is OK but can also be confusing when a user is trying to figure out how to return the right thing. Maybe there should be an opt-in strict mode (default off).

Hmm ... my initial reaction is that this is not a bad idea at all.
Are you proposing something like this?
```python
out = array.map_blocks(partial(f, ndim=8), chunks=out_chunks)
with dask.config.set({"strict_chunk_ndim_check": True}):
    arr2 = out.compute()  # raises an informative error regarding chunk dimensions
```

This could also be helpful for debugging.  I'll certainly give it some more thought.
Note that I've made changes to this PR so that, @dcherian 's original problem, for example, now raises a `ValueError` with a more user-friendly error message within the context of a new `dask` configuration setting "array.computing.check-chunk-ndim", otherwise `dask` does its thing and no exception is raised.  Is this satisfactory?

```python
import numpy as np
import dask.array as da


array = da.from_array([1, 1, 2, 2, 2, 3, 3, 3, 3, 3], chunks=((2, 3, 5),))
output_chunks = ((1, 1, 1),)

with dask.config.set({"array.computing.check-chunk-ndim": True}):
    array.map_blocks(np.mean, chunks=output_chunks).compute()  # raises Dimension mismatch error
```
Error trace log:

    Traceback (most recent call last):
      File "C:\Users\heamu\git\dask\dask\array\tests\dcherian.py", line 11, in <module>
        array.map_blocks(np.mean, chunks=output_chunks).compute()
      File "C:\Users\heamu\git\dask\dask\base.py", line 292, in compute
        (result,) = compute(self, traverse=False, **kwargs)
      File "C:\Users\heamu\git\dask\dask\base.py", line 576, in compute
        return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
      File "C:\Users\heamu\git\dask\dask\base.py", line 576, in <listcomp>
        return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
      File "C:\Users\heamu\git\dask\dask\array\core.py", line 1215, in finalize
        return concatenate3(results)
      File "C:\Users\heamu\git\dask\dask\array\core.py", line 5091, in concatenate3
        chunks = chunks_from_arrays(arrays, ndim)
      File "C:\Users\heamu\git\dask\dask\array\core.py", line 4878, in chunks_from_arrays
        result.append(tuple(shape(deepfirst(a), expected_ndim)[dim] for a in arrays))
      File "C:\Users\heamu\git\dask\dask\array\core.py", line 4878, in <genexpr>
        result.append(tuple(shape(deepfirst(a), expected_ndim)[dim] for a in arrays))
      File "C:\Users\heamu\git\dask\dask\array\core.py", line 4870, in shape
        raise ValueError(
    ValueError: Dimension mismatch: expected chunk ndim = 1. Got instead ndim = 0.

I'm unsure whether using a `dask` configuration setting is the right way to obtain this behavior though.  An alternative way might be to provide a new keyword argument to `dask.array.Array.compute()`.

What do y'all think?
Got it.

I've removed the configuration option and instead added a keyword argument `check_ndim` for `.compute()` instead.  

However, I could not make `check_ndim=True` the default without breaking existing tests.  Looks like `dask` itself depends on `check_ndim=False` after all.

So under this PR we now have:
```python
import numpy as np
import dask.array as da


array = da.from_array([1, 1, 2, 2, 2, 3, 3, 3, 3, 3], chunks=((2, 3, 5),))
output_chunks = ((1, 1, 1),)

array = array.map_blocks(np.mean, chunks=output_chunks).persist()
array.compute(check_ndim=True)  # raises "ValueError: Dimension mismatch: expected chunk ndim = 1. Got instead ndim = 0."
```

The other option is to make `assert_eq(..., check_chunks=True)` stricter:
https://github.com/dask/dask/blob/1b1a96b4ddf31d0879e718be463a731af7f1d086/dask/array/utils.py#L208-L214

which is called by https://github.com/dask/dask/blob/1b1a96b4ddf31d0879e718be463a731af7f1d086/dask/array/utils.py#L226

That check passes with (unexpectedly to me)
```
ipdb>  expected_shape
(1,)

ipdb>  chunk.shape
(1, 1, 1, 1, 1, 1, 1, 1)
```
@dcherian 

Good idea.  I've just introduced a new keyword argument `check_ndim`  for `assert_eq()` which is `True` by default.  I just hope it doesn't break any tests in `dask`.
Seems to have worked!
Thanks for the PR, `@ParticularMiner`! 

@jsignell Thank you for reviewing this! It looks complete, are we good to merge?
Sorry I dropped the ball on this one. Let me look it over one last time.
Sorry for the slowness here. Looking at this with fresh eyes, I don't think this merits a top-level compute kwarg. I know my earlier comments were unclear, but I was imagining a `map_blocks` kwarg, not a `compute` kwarg.

That being said I understand the motivation to abstract this to the compute level and if you geel strongly that that is where it goes I definitely prefer the config option approach that you had a few commits back. 
No worries @jsignell !

The way I see it, a `map_blocks()` `check_ndims` kwarg would have to check the chunk-function output at graph-construction time, which in my opinion goes against `dask`'s modus operandi of delaying the execution of chunk-functions.

Alternatively, `map_blocks(..., check_ndims=True)` could "append" a new built-in chunk function, say `check_chunk_ndim()`, to the task-graph of the output `dask` array which will be then be executed at graph-execution time.   (Perhaps, this option is what you're leaning more towards?) 
> The way I see it, a `map_blocks()` `check_ndims` kwarg would have to check the chunk-function output at graph-construction time, which in my opinion goes against `dask`'s modus operandi of delaying the execution of chunk-functions.

Yes that is definitely not a good idea :)

> Alternatively, `map_blocks(..., check_ndims=True)` could "append" a new built-in chunk function, say `check_chunk_ndim()`, to the task-graph of the output `dask` array which will be then be executed at graph-execution time. (Perhaps, this option is what you're leaning more towards?)

I had not fully thought through the implementation but I was imagining that this would work roughly how `enforce_metadata` does on the dask.dataframe side. Basically it wraps the user function and applies some validation after the function is computed.

https://github.com/dask/dask/blob/2115574642af273c244ac4fc372ada9d6e69f931/dask/dataframe/core.py#L6405-L6414


@jsignell 

I think I get it now.  Indeed, your suggestion is similar to the second option I suggested.   I'll see what I can do.
@jsignell 

Please see the new diffs.  

As promised, I've added a kwarg `enforce_ndim` which works in a similar way to `enforce_metadata` for `dask.dataframe.map_partitions()`.
@jsignell 

Looks like all tests have passed!

I think your suggestion is much better than either a configuration setting or a `compute()` kwarg.  

I've also retained @dcherian 's suggestion for `assert_eq(, check_ndim=True)`.

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -556,11 +556,11 @@ def map_blocks(
     new_axis : number or iterable, optional
         New dimensions created by the function. Note that these are applied
         after ``drop_axis`` (if present).
-    enforce_ndim : bool, default True
+    enforce_ndim : bool, default False
         Whether to enforce at runtime that the dimensionality of the array
         produced by ``func`` actually matches that of the array returned by
         ``map_blocks``.
-        This will raise an error if there is a mismatch.
+        If True, this will raise an error when there is a mismatch.
     token : string, optional
         The key prefix to use for the output array. If not provided, will be
         determined from the function name.


Test patch: