Update apply_ufunc output_sizes error message (#7509)

* test for error message

* fix

* whatsnew

---------

Co-authored-by: Deepak Cherian <[email protected]>

Author: TomNicholas

doc/whats-new.rst

@@ -58,6 +58,8 @@ Bug fixes
   By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_ and `Scott Chamberlin <https://github.com/scottcha>`_.
 - Handle ``keep_attrs`` option in binary operators of :py:meth:`Dataset` (:issue:`7390`, :pull:`7391`).
   By `Aron Gergely <https://github.com/arongergely>`_.
+- Improve error message when using dask in :py:func:`apply_ufunc` with ``output_sizes`` not supplied. (:pull:`7509`)
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 - :py:func:`xarray.Dataset.to_zarr` now drops variable encodings that have been added by xarray during reading
   a dataset. (:issue:`7129`, :pull:`7500`).
   By `Hauke Schulz <https://github.com/observingClouds>`_.

xarray/core/computation.py

@@ -723,7 +723,9 @@ def apply_variable_ufunc(
                 dask_gufunc_kwargs["output_sizes"] = output_sizes_renamed
 
             for key in signature.all_output_core_dims:
-                if key not in signature.all_input_core_dims and key not in output_sizes:
+                if (
+                    key not in signature.all_input_core_dims or key in exclude_dims
+                ) and key not in output_sizes:
                     raise ValueError(
                         f"dimension '{key}' in 'output_core_dims' needs corresponding (dim, size) in 'output_sizes'"
                     )

xarray/tests/test_computation.py

@@ -1220,6 +1220,25 @@ def func(da):
     assert_identical(expected.chunk(), actual)
 
 
+@requires_dask
+def test_apply_dask_new_output_sizes_not_supplied_same_dim_names() -> None:
+    # test for missing output_sizes kwarg sneaking through
+    # see GH discussion 7503
+
+    data = np.random.randn(4, 4, 3, 2)
+    da = xr.DataArray(data=data, dims=("x", "y", "i", "j")).chunk(x=1, y=1)
+
+    with pytest.raises(ValueError, match="output_sizes"):
+        xr.apply_ufunc(
+            np.linalg.pinv,
+            da,
+            input_core_dims=[["i", "j"]],
+            output_core_dims=[["i", "j"]],
+            exclude_dims=set(("i", "j")),
+            dask="parallelized",
+        )
+
+
 def pandas_median(x):
     return pd.Series(x).median()