Repository: pydata/xarray (Python)

Issue:
Set `allow_rechunk=True` in `apply_ufunc`
<!-- 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**:

`blockwise` calls `unify_chunks` by default but `apply_gufunc` does not; so we have a regression in `apply_ufunc` now that we've switched from `blockwise` to `apply_gufunc`.

**Minimal Complete Verifiable Example**:

```python
import operator

a = xr.DataArray(np.arange(10), dims=("a")).chunk({"a": 2})
b = xr.DataArray(np.arange(10), dims=("a")).chunk({"a": 4})

xr.apply_ufunc(operator.add, a, b, dask="parallelized", output_dtypes=[a.dtype]).compute()
```

raises
```
ValueError: Dimension `'__loopdim0__'` with different chunksize present
```

on master but works with 0.16.0

I think we need to do `dask_gufunc_kwargs.setdefault("allow_rechunk", True)`

If we want to avoid that, we'll need to go through a deprecation cycle.


Discussion:
(this is causing downstream test failures: https://github.com/NCAR/pop-tools/pull/59; thanks @mnlevy1981)
The behaviour changed in #4060 (commit https://github.com/pydata/xarray/commit/a7fb5a9fa1a2b829181ea9e4986b959f315350dd). Please see discussion with regard to `allow_rechunk` over there. Reason to not handle/set `allow_rechunk=True` was @shoyer's [comment](https://github.com/pydata/xarray/pull/4060#issuecomment-634776667).
Copying over that comment...

> I am pretty confident that existing behavior of xarray.apply_ufunc with un-equal chunks along non-core dimensions is entirely broken. I am OK with just erroring for now.

@shoyer could you please clarify what you meant?

For example, this works with v0.16.0 but fails on master
``` python
import operator

import numpy as np

import xarray as xr

a = xr.DataArray(np.ones((10, 10)), dims=("a", "b")).chunk({"a": 2, "b": 1})
b = xr.DataArray(np.ones((10, 10)), dims=("a", "b")).chunk({"a": -1, "b": 4})

xr.apply_ufunc(
    operator.add, a, b, dask="parallelized", output_dtypes=[a.dtype],
).compute().equals(a.compute() + b.compute())
```
One solution would be to catch this ValueError, issue a FutureWarning and add `allow_rechunk=True` to `dask_gufunc_kwargs` here:

https://github.com/pydata/xarray/blob/9c85dd5f792805bea319f01f08ee51b83bde0f3b/xarray/core/computation.py#L646-L657

```python
        def func(*arrays):
            import dask.array as da

            gufunc = functools.partial(
                da.apply_gufunc,
                numpy_func,
                signature.to_gufunc_string(exclude_dims),
                *arrays,
                vectorize=vectorize,
                output_dtypes=output_dtypes,
            )

            try:
                res = gufunc(**dask_gufunc_kwargs)
            except ValueError as exc:
                if "with different chunksize present" in str(exc):
                    warnings.warn(
                        f"``allow_rechunk=True`` need to be explicitely set in the "
                        f"``dask_gufunc_kwargs`` parameter. Not setting will raise dask "
                        f"ValueError ``{str(exc)}`` in a future version.",
                        FutureWarning,
                        stacklevel=2,
                    )
                    dask_gufunc_kwargs["allow_rechunk"] = True
                    res = gufunc(**dask_gufunc_kwargs)
                else:
                    raise
```

I could make a PR out of this. The message wording can surely be improved. WDYT @dcherian and @shoyer?

Maybe we do want to set allow_rechunk=True? It seems that I was just mistaken about the current behavior.
@shoyer In this case: Should we warn the user, that data might be loaded into memory?

Another questions are, why does this kwarg exist in dask and why do they not rechunk per default? 
From the dask `apply_gufunc` docstring:

```python
"""
allow_rechunk: Optional, bool, keyword only

    Allows rechunking, otherwise chunk sizes need to match and core dimensions are to consist only of one chunk. 
    Warning: enabling this can increase memory usage significantly. Defaults to False
"""
```
Current code handling in dask:

https://github.com/dask/dask/blob/42873f27ce11ce35652dda344dae5c47b742bef2/dask/array/gufunc.py#L398-L417

```python
         if not allow_rechunk:
            chunksizes = chunksizess[dim]
            #### Check if core dimensions consist of only one chunk
            if (dim in core_shapes) and (chunksizes[0][0] < core_shapes[dim]):
                raise ValueError(
                    "Core dimension `'{}'` consists of multiple chunks. To fix, rechunk into a single \
chunk along this dimension or set `allow_rechunk=True`, but beware that this may increase memory usage \
significantly.".format(
                        dim
                    )
                )
            #### Check if loop dimensions consist of same chunksizes, when they have sizes > 1
            relevant_chunksizes = list(
                unique(c for s, c in zip(sizes, chunksizes) if s > 1)
            )
            if len(relevant_chunksizes) > 1:
                raise ValueError(
                    "Dimension `'{}'` with different chunksize present".format(dim)
                )
```

IIUTC, this not only rechunks non-core dimensions but also fixes core dimensions with more than one chunk. Would  this be intended from the `xarray-side`? Before #4060 core dimension chunks>1 was catched and errored:

```python
            # core dimensions cannot span multiple chunks
            for axis, dim in enumerate(core_dims, start=-len(core_dims)):
                if len(data.chunks[axis]) != 1:
                    raise ValueError(
                        "dimension {!r} on {}th function argument to "
                        "apply_ufunc with dask='parallelized' consists of "
                        "multiple chunks, but is also a core dimension. To "
                        "fix, rechunk into a single dask array chunk along "
                        "this dimension, i.e., ``.chunk({})``, but beware "
                        "that this may significantly increase memory usage.".format(
                            dim, n, {dim: -1}
                        )
                    )
```

Explicit `rechunk` was recommended to the user, though.

That means setting `allow_rechunk=True` per default alone will not give us same behaviour as before #4060. I'm unsure how to proceed.
> Another questions are, why does this kwarg exist in dask and why do they not rechunk per default?

Trying to answer this from looking at the dask code. 

- `allow_rechunk=False`: catch chunking problems in core and non-core-dimensions and raise an error.
   This helps to prevent users running into loading huge dask arrays into memory without further notice.
- `allow_rechunk=True`: blockwise is called with `align_arrays=True` per default which means automatic rechunking for all arrays (core and non-core dimensions). Users can use this, if they are sure the system can handle possible large amounts of data.

So to maintain backward compatibility, we should add that same check
``` python
            # core dimensions cannot span multiple chunks
            for axis, dim in enumerate(core_dims, start=-len(core_dims)):
                if len(data.chunks[axis]) != 1:
                    raise ValueError(
                        "dimension {!r} on {}th function argument to "
                        "apply_ufunc with dask='parallelized' consists of "
                        "multiple chunks, but is also a core dimension. To "
                        "fix, rechunk into a single dask array chunk along "
                        "this dimension, i.e., ``.chunk({})``, but beware "
                        "that this may significantly increase memory usage.".format(
                            dim, n, {dim: -1}
                        )
                    )
```

and set `allow_rechunk=True`. 

We could deprecate and remove this check in a couple of versions but I don't know if it's worth the effort...
@dcherian @shoyer 

In #4392 I've tried to get around this bug. I found it easier to just catch the dask ValueError's and not add more code checks. I'll add more information in that PR.



Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index eda5c9dfbf6..0a9a1464df1 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -96,7 +96,7 @@ Documentation
 Internal Changes
 ~~~~~~~~~~~~~~~~
 - Use :py:func:`dask.array.apply_gufunc` instead of :py:func:`dask.array.blockwise` in
-  :py:func:`xarray.apply_ufunc` when using ``dask='parallelized'``. (:pull:`4060`, :pull:`4391`)
+  :py:func:`xarray.apply_ufunc` when using ``dask='parallelized'``. (:pull:`4060`, :pull:`4391`, :pull:`4392`)
 - Fix ``pip install .`` when no ``.git`` directory exists; namely when the xarray source
   directory has been rsync'ed by PyCharm Professional for a remote deployment over SSH.
   By `Guido Imperiale <https://github.com/crusaderky>`_
diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index 507f12fe55e..db4a9b46e98 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -625,6 +625,26 @@ def apply_variable_ufunc(
             if dask_gufunc_kwargs is None:
                 dask_gufunc_kwargs = {}
 
+            allow_rechunk = dask_gufunc_kwargs.get("allow_rechunk", None)
+            if allow_rechunk is None:
+                for n, (data, core_dims) in enumerate(
+                    zip(input_data, signature.input_core_dims)
+                ):
+                    if is_duck_dask_array(data):
+                        # core dimensions cannot span multiple chunks
+                        for axis, dim in enumerate(core_dims, start=-len(core_dims)):
+                            if len(data.chunks[axis]) != 1:
+                                raise ValueError(
+                                    f"dimension {dim} on {n}th function argument to "
+                                    "apply_ufunc with dask='parallelized' consists of "
+                                    "multiple chunks, but is also a core dimension. To "
+                                    "fix, either rechunk into a single dask array chunk along "
+                                    f"this dimension, i.e., ``.chunk({dim}: -1)``, or "
+                                    "pass ``allow_rechunk=True`` in ``dask_gufunc_kwargs`` "
+                                    "but beware that this may significantly increase memory usage."
+                                )
+                dask_gufunc_kwargs["allow_rechunk"] = True
+
             output_sizes = dask_gufunc_kwargs.pop("output_sizes", {})
             if output_sizes:
                 output_sizes_renamed = {}


Test patch:
diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
index 7cb755b6dac..63bedfaf280 100644
--- a/xarray/tests/test_computation.py
+++ b/xarray/tests/test_computation.py
@@ -695,8 +695,7 @@ def check(x, y):
     check(data_array, 0 * data_array)
     check(data_array, 0 * data_array[0])
     check(data_array[:, 0], 0 * data_array[0])
-    with raises_regex(ValueError, "with different chunksize present"):
-        check(data_array, 0 * data_array.compute())
+    check(data_array, 0 * data_array.compute())
 
 
 @requires_dask
@@ -710,7 +709,7 @@ def test_apply_dask_parallelized_errors():
     with raises_regex(ValueError, "at least one input is an xarray object"):
         apply_ufunc(identity, array, dask="parallelized")
 
-    # formerly from _apply_blockwise, now from dask.array.apply_gufunc
+    # formerly from _apply_blockwise, now from apply_variable_ufunc
     with raises_regex(ValueError, "consists of multiple chunks"):
         apply_ufunc(
             identity,