Took a look and it seems to originate from the stacking part and someting in dask.

In polyfit, we rearrange the DataArrays to 2D arrays, so we can run the least squares with np/dsa.apply_along_axis. But I checked and the chunking problem seems to appear before any call of the sort. MWE:

```python3 import xarray as xr import dask.array as dsa nz, ny, nx = (10, 20, 30) data = dsa.ones((nz, ny, nx), chunks=(1, 5, nx)) da = xr.DataArray(data, dims=['z', 'y', 'x']) da.chunks

((1, 1, 1, 1, 1, 1, 1, 1, 1, 1), (5, 5, 5, 5), (30,))

stk = da.stack(zy=['z', 'y']) print(stk.dims, stk.chunks)

('x', 'zy') ((30,), (20, 20, 20, 20, 20, 20, 20, 20, 20, 20))

Merged chunks!

```

And then I went down the rabbit hole (ok it's not that deep) and is all goes down here: https://github.com/pydata/xarray/blob/e0358e586079c12525ce60c4a51b591dc280713b/xarray/core/variable.py#L1507

In Variable._stack_one the stacking is performed and Variable.data.reshape is called. Dask itself is rechunking the output, merging the chunks. There is a merge_chunks kwarg for reshape, but I think it has a bug:

```python

Let's stack as xarray does: x, z, y -> x, zy

data_t = data.transpose(2, 0, 1) # Dask array with shape (30, 10, 20), the same as reordered in Variable._stack_once.

new_data = data_t.reshape((30, -1), merge_chunks=True) # True is the default, this is the same call as in xarray new_data.chunks

((30,), (20, 20, 20, 20, 20, 20, 20, 20, 20, 20))

new_data = data_t.reshape((30, -1), merge_chunks=False) new_data.shape # I'm printing shape because chunks is too large, but see the bug:

(30, 6000) # instead of (30, 200)!!!

Doesn't happen when we do not transpose. So let's reshape data as z, y, x -> zy, x

new_data = data.reshape((-1, 30), merge_chunks=True) new_data.chunks

((5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5), (30,))

Chunks were not merged? But this is the output expected by paigem.

new_data = data.reshape((-1, 30), merge_chunks=False) new_data.chunks

((5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5), (30,))

That's what I expected with merge_chunks=False.

```

For polyfit itself, the apply_along_axis call could be changed to a apply_ufunc with vectorize=True, I think this would avoid the problem and behave the same on the user's side. Would need some refactoring.