Repository: dask/dask

Issue:
Unwanted aggregation `map_blocks`
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:

While working on https://github.com/dask/dask/issues/8822 I tried to use `drop_axis=True` in `map_blocks` and ran into an issue where `map_blocks` aggregated the results from each block rather than concatenating them.

**What you expected to happen**:

The output should contain one output per block. 

```python
array([1., 2., 3.])
```

**Minimal Complete Verifiable Example**:

```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)))
array.map_blocks(np.mean, drop_axis=[0]).compute()
# 2.3
```

Ping @rjzamora for blockwise insight. 



Discussion:
I'm not as familiar with `dask.array` as I'd like to be, but doesn't passing `drop_axis=[0]` to `map_blocks` mean you want to concatenate along the only axis you have?
Hmm I guess I thought it would mean "the output of my chunk-function will have a different dimensionality than the input" But I still expect one result per block.
@jsignell @rjzamora 

The [`blockwise()` documentation](https://docs.dask.org/en/latest/generated/dask.array.blockwise.html#dask.array.blockwise) states:

> Any index, like `i` missing from the output index is interpreted as a contraction ... the passed function should expect an iterable of blocks on any array that holds that index. To receive arrays concatenated along contracted dimensions instead pass concatenate=True.

So `map_blocks()` calls `blockwise()` without any output axis labels (due to `drop_axis=[0]`).  But `blockwise()` compares input and output axis labels and sees that the output  is missing the only axis label that the input has so it concatenates all chunks of the input into one large chunk which `np.mean` then reduces to a scalar.

As far as I can see, this behavior should be expected as this is the only way for the return value of `map_blocks()`/`blockwise()` to remain consistent with the expected output dimensionality, or meta data.  I mean, if the output dask array is expected to have a dimensionality of 0, then only a single value must be returned, therefore the dask array must have only one chunk containing that value.  The only way to enforce this is to concatenate all the input chunks into one chunk before passing it to the chunk-function.
I'm not sure I agree with the "expected result" of
```python
array([1., 2., 3.])
```
To me, the fundamental contract of `map_blocks` is that the function should transform an array-like into an array-like. There might be changes to the representation (like sparsity or cupy or whatever), or the shape, but it should always be an arraylike, and  returning a scalar violates that contract. Automatically concatenating the resulting scalars into a new array would have some pretty surprising and confusing results that, to me, should be avoided.

To most users, I think we should be striving to make chunks an implementation detail. Clearly, the performance characteristics of a computation will depend on chunking, but I don't think the answer should. If my code got different answers depending on how I chose my chunked, I'd be very puzzled, and it would be pretty easy to hit that scenario if I simply forgot to specify which concatenation axis I want for `np.mean`. So with your example above

```
array = dask.array.from_array([1,1,2,2,2,3,3,3,3,3], chunks=((2,3,5)))
array.map_blocks(np.mean, drop_axis=[0]).compute()  # [1, 2, 3]


array = dask.array.from_array([1,1,2,2,2,3,3,3,3,3], chunks=((5,5)))
array.map_blocks(np.mean, drop_axis=[0]).compute()  # [1.6, 3]  Weird! 
```

Now, users could still make cases that have that above behavior by ensuring that their reduction function keeps their axes (this gets your expected result):

```python
array = dask.array.from_array([1,1,2,2,2,3,3,3,3,3], chunks=((2,3,5)))
array.map_blocks(np.mean, keepdims=True).compute()  # [1, 2, 3]
```

But I guess I'd be in favor of *making* people opt-in to chunking-dependent behavior with array-forcing workarounds like the above. So I guess I agree with @dcherian's suggestion in #8822 that a better error message if a scalar is returned is the right behavior.
Hmm yeah I can see how the behavior that I am describing in this issue might be unwanted and surprising. I guess I'm more familiar with the dataframe equivalent which _does_ always return at least one value per partition and wraps scalars in Series objects if needed. But it is entirely possible that they are different enough that it's not worth trying to make their behavior internally consistent. 

The name `map_blocks` to me implies a simple mapping, so the fact that it's possible to get a cross-block aggregation is what really confused me. I understand how it's possible given that `map_blocks` is backed by `blockwise` but I'm not sure it _should_ be possible. 

But this is probably orthogonal to @dcherian's original issue. I totally agree that an error message makes sense in that case. It can probably be raised at graph construction.

Yeah, the documentation is lacking on the precise nature of the mapping underlying `map_blocks()` when its input parameter `drop_axis` is not `None`.  It is, in general, a "many-blocks-to-one-block" mapping under such a condition.  An analogous mapping in this case is the projection from, say, 3d to 2d space:

(*x*, *y*, *z*) &rightarrow; (*x*, *y*).

One could even think of it as *somewhat* like a reverse of broadcasting.

And I must admit that at first it surprised me too, until I looked at the issue more closely and realized that it made sense and that I probably shouldn't have been surprised.  After all, when an entire axis is dropped, where should all the blocks along that axis be mapped to? 

So perhaps, while @dcherian's issue is being fixed, a brief note regarding this could be added to the docs to alert users.
Here's an illustrative diagram that could be added to the docs to clarify the case of `drop_axis != None`:

![image](https://user-images.githubusercontent.com/78448465/160272226-965e8733-e958-4263-b70d-ae4f2c7722dd.png)

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -618,6 +618,12 @@ def map_blocks(
     If ``chunks`` is specified but ``new_axis`` is not, then it is inferred to
     add the necessary number of axes on the left.
 
+    Note that ``map_blocks()`` will concatenate chunks along axes specified by
+    the keyword parameter ``drop_axis`` prior to applying the function.
+    This is illustrated in the figure below:
+
+    .. image:: /images/map_blocks_drop_axis.png
+
     Map_blocks aligns blocks by block positions without regard to shape. In the
     following example we have two arrays with the same number of blocks but
     with different shape and chunk sizes.


Test patch: