Repository: dask/dask

Issue:
support meta in apply_gufunc
<!-- 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**:

`dask.array.apply_ufunc` does not accept a `meta` kwarg.

**Motivation**

There is a PR (https://github.com/pydata/xarray/pull/4060) updating `xarray.apply_ufunc` to use `dask.array.apply_gufunc`. Currently `apply_ufunc` used `blockwise` which accepted a `meta` kwarg but now we see that `apply_gufunc` does not accept one.

I don't have an example that fails with the current code. But it seems like a good idea to let the user specify `meta` in case the automated inference in the following lines doesn't work
https://github.com/dask/dask/blob/ad5006fc70a256ceb389a3e42b6f303c61c1721d/dask/array/gufunc.py#L420-L440


Discussion:
IIUC @shoyer's [comment in 4060](https://github.com/pydata/xarray/pull/4060#issuecomment-634774468) `meta` can still be forwarded into `apply_gufunc` via  `**kwargs`. And this is what is done in the current state of https://github.com/pydata/xarray/pull/4060. 

But looking at the above code, this would break if the except-clause will be reached (`meta` inside `**kwargs`). So, explicitely supporting `meta` kwarg would make sense.
Agreed, accepting `meta` here makes since. Is anyone interested in working on this?

Resolving patch:
diff --git a/dask/array/gufunc.py b/dask/array/gufunc.py
--- a/dask/array/gufunc.py
+++ b/dask/array/gufunc.py
@@ -240,6 +240,9 @@ def apply_gufunc(func, signature, *args, **kwargs):
         dimensions are to consist only of one chunk.
         Warning: enabling this can increase memory usage significantly.
         Defaults to ``False``.
+    meta: Optional, tuple, keyword only
+        tuple of empty ndarrays describing the shape and dtype of the output of the gufunc.
+        Defaults to ``None``.
     **kwargs : dict
         Extra keyword arguments to pass to `func`
 
@@ -279,6 +282,7 @@ def apply_gufunc(func, signature, *args, **kwargs):
     output_sizes = kwargs.pop("output_sizes", None)
     vectorize = kwargs.pop("vectorize", None)
     allow_rechunk = kwargs.pop("allow_rechunk", False)
+    meta = kwargs.pop("meta", None)
 
     # Input processing:
     ## Signature
@@ -417,27 +421,32 @@ def apply_gufunc(func, signature, *args, **kwargs):
     ### Use existing `blockwise` but only with loopdims to enforce
     ### concatenation for coredims that appear also at the output
     ### Modifying `blockwise` could improve things here.
-    try:
-        tmp = blockwise(  # First try to compute meta
-            func, loop_output_dims, *arginds, concatenate=True, **kwargs
-        )
-    except ValueError:
-        # If computing meta doesn't work, provide it explicitly based on
-        # provided dtypes
-        sample = arginds[0]._meta
-        if isinstance(output_dtypes, tuple):
-            meta = tuple(
-                meta_from_array(sample, dtype=odt)
-                for ocd, odt in zip(output_coredimss, output_dtypes)
-            )
-        else:
-            meta = tuple(
-                meta_from_array(sample, dtype=odt)
-                for ocd, odt in zip((output_coredimss,), (output_dtypes,))
-            )
+    if meta is not None:
         tmp = blockwise(
             func, loop_output_dims, *arginds, concatenate=True, meta=meta, **kwargs
         )
+    else:
+        try:
+            tmp = blockwise(  # First try to compute meta
+                func, loop_output_dims, *arginds, concatenate=True, **kwargs
+            )
+        except ValueError:
+            # If computing meta doesn't work, provide it explicitly based on
+            # provided dtypes
+            sample = arginds[0]._meta
+            if isinstance(output_dtypes, tuple):
+                meta = tuple(
+                    meta_from_array(sample, dtype=odt)
+                    for ocd, odt in zip(output_coredimss, output_dtypes)
+                )
+            else:
+                meta = tuple(
+                    meta_from_array(sample, dtype=odt)
+                    for ocd, odt in zip((output_coredimss,), (output_dtypes,))
+                )
+            tmp = blockwise(
+                func, loop_output_dims, *arginds, concatenate=True, meta=meta, **kwargs
+            )
 
     if isinstance(tmp._meta, tuple):
         metas = tmp._meta
@@ -628,7 +637,7 @@ def __call__(self, *args, **kwargs):
             output_sizes=self.output_sizes,
             output_dtypes=self.output_dtypes,
             allow_rechunk=self.allow_rechunk or kwargs.pop("allow_rechunk", False),
-            **kwargs
+            **kwargs,
         )
 
 
@@ -681,6 +690,9 @@ def as_gufunc(signature=None, **kwargs):
         dimensions are to consist only of one chunk.
         Warning: enabling this can increase memory usage significantly.
         Defaults to ``False``.
+    meta: Optional, tuple, keyword only
+        tuple of empty ndarrays describing the shape and dtype of the output of the gufunc.
+        Defaults to ``None``.
 
     Returns
     -------
@@ -720,6 +732,7 @@ def as_gufunc(signature=None, **kwargs):
         "output_sizes",
         "output_dtypes",
         "allow_rechunk",
+        "meta",
     }
     if set(_allowedkeys).issubset(kwargs.keys()):
         raise TypeError("Unsupported keyword argument(s) provided")


Test patch:
diff --git a/dask/array/tests/test_gufunc.py b/dask/array/tests/test_gufunc.py
--- a/dask/array/tests/test_gufunc.py
+++ b/dask/array/tests/test_gufunc.py
@@ -174,6 +174,7 @@ def outer_product(x, y):
     a = da.random.normal(size=(20, 30), chunks=(5, 30))
     b = da.random.normal(size=(10, 1, 40), chunks=(10, 1, 40))
     c = apply_gufunc(outer_product, "(i),(j)->(i,j)", a, b, output_dtypes=a.dtype)
+
     assert c.compute().shape == (10, 20, 30, 40)
 
 
@@ -615,3 +616,15 @@ def stats(x):
 
     assert_eq(sum, sum)
     assert_eq(mean, mean)
+
+
+def test_apply_gufunc_with_meta():
+    def stats(x):
+        return np.mean(x, axis=-1), np.std(x, axis=-1, dtype=np.float32)
+
+    a = da.random.normal(size=(10, 20, 30), chunks=(5, 5, 30))
+    meta = (np.ones(0, dtype=np.float64), np.ones(0, dtype=np.float32))
+    result = apply_gufunc(stats, "(i)->(),()", a, meta=meta)
+    expected = stats(a.compute())
+    assert_eq(expected[0], result[0])
+    assert_eq(expected[1], result[1])