Repository: dask/dask
Issue category: Bug Report

Issue:
`map_blocks()` with `new_axis` output has incorrect shape
<!-- 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.
-->

**Describe the issue**:
When running `map_blocks()` with `new_axis` specified, the output shape of the dask array is not set correctly. In the below example I would expect it to be the same as the shape after computation.

**Minimal Complete Verifiable Example**:

```python
import dask.array as da
import numpy as np

def func(x):
    return np.stack([x, x + 0.5])

x = da.arange(6, chunks=2)
x_mapped = x.map_blocks(func, new_axis=[0])

print(x_mapped.shape)
# (1, 6)
print(x_mapped.compute().shape)
# (2, 6)
```

**Anything else we need to know?**:

**Environment**:

- Dask version: b7d9bf49f682de8d2ef51f4617e3da782400c290
- Python version: 3.12.3
- Operating System: macOS
- Install method (conda, pip, source): source



Hints:
I don't think that we can guess the output shape with a high degree of fidelity. We should probably either set all chunks to NaN or force the specification of chunks.
Being able to specify the size of new output dimensions if known would be nice. e.g., in the above toy example we know the size of the new dimension is going to be `2` ahead of time.
Yeah, the interesting thing for Dask is the chunk sizes, the shape is just a convenient result from that, so we would need this information.

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index fdf65bd24a4..3065406a922 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -562,7 +562,9 @@ def map_blocks(
         Dimensions lost by the function.
     new_axis : number or iterable, optional
         New dimensions created by the function. Note that these are applied
-        after ``drop_axis`` (if present).
+        after ``drop_axis`` (if present). The size of each chunk along this
+        dimension will be set to 1. Please specify ``chunks`` if the individual
+        chunks have a different size.
     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


Test patch: