Repository: dask/dask

Issue:
`apply_gufunc` bugs with meta inference
**What happened**:
`apply_gufunc` returns a tuple with the wrong number of outputs for multi-output signatures.

**What you expected to happen**:
`apply_gufunc` would return the same number of items as the signature implied.

**Minimal Complete Verifiable Example**:

```python
import dask.array as da
from dask.array import apply_gufunc
import numpy

>>> len(apply_gufunc(lambda a, b: (a.max(axis=1), b.max()), '(i,I),(I)->(i),()', da.ones((10, 10), chunks=-1), da.ones(10, chunks=-1), output_dtypes=['float64', 'float64']))
1
```
This should return a tuple of 2 items (one array and one scalar), as the signature implies.

The issue is that `meta_from_array` encounters a "zero-size array to reduction operation" error. This error used to propagate (which would have taken us down a [more useful codepath](https://github.com/dask/dask/blob/ac1bd05cfd40207d68f6eb8603178d7ac0ded922/dask/array/gufunc.py#L432-L448) in `apply_gufunc`), but as of recently (#6736) it just selects the first element from `args_meta`: https://github.com/dask/dask/blob/ac1bd05cfd40207d68f6eb8603178d7ac0ded922/dask/array/utils.py#L158-L161

@pentschev I wonder if this `meta = args_meta[0]` should be conditioned on `len(args_meta) = 1`? When there are multiple arguments, I think the only sensible behavior is to return None, since we don't know how the arguments will be combined. Even when there's only a single argument, I personally thing it's incorrect to assume in generality that the output type will be the same as the input type; an arbitrary user-defined function could do all sorts of other things before/after calling `np.min`.

Even if `meta_from_array` had raised an error, this still wouldn't have worked. This logic for generating meta from the `output_dtypes` https://github.com/dask/dask/blob/ac1bd05cfd40207d68f6eb8603178d7ac0ded922/dask/array/gufunc.py#L436-L440 will rarely/never be reached, because `output_dtypes` is typically a list, not a tuple.

Overall, there is inconsistency throughout `apply_gufunc` as to whether `nout`, the type+length of `output_dtypes`, or the type+length of `metas` is the source of truth for how many items to return.


**Anything else we need to know?**:

**Environment**:

- Dask version: ac1bd05cfd40207d68f6eb8603178d7ac0ded922
- Python version: 3.8.8
- Operating System: macOS
- Install method (conda, pip, source): source

`apply_gufunc` bugs with meta inference
**What happened**:
`apply_gufunc` returns a tuple with the wrong number of outputs for multi-output signatures.

**What you expected to happen**:
`apply_gufunc` would return the same number of items as the signature implied.

**Minimal Complete Verifiable Example**:

```python
import dask.array as da
from dask.array import apply_gufunc
import numpy

>>> len(apply_gufunc(lambda a, b: (a.max(axis=1), b.max()), '(i,I),(I)->(i),()', da.ones((10, 10), chunks=-1), da.ones(10, chunks=-1), output_dtypes=['float64', 'float64']))
1
```
This should return a tuple of 2 items (one array and one scalar), as the signature implies.

The issue is that `meta_from_array` encounters a "zero-size array to reduction operation" error. This error used to propagate (which would have taken us down a [more useful codepath](https://github.com/dask/dask/blob/ac1bd05cfd40207d68f6eb8603178d7ac0ded922/dask/array/gufunc.py#L432-L448) in `apply_gufunc`), but as of recently (#6736) it just selects the first element from `args_meta`: https://github.com/dask/dask/blob/ac1bd05cfd40207d68f6eb8603178d7ac0ded922/dask/array/utils.py#L158-L161

@pentschev I wonder if this `meta = args_meta[0]` should be conditioned on `len(args_meta) = 1`? When there are multiple arguments, I think the only sensible behavior is to return None, since we don't know how the arguments will be combined. Even when there's only a single argument, I personally thing it's incorrect to assume in generality that the output type will be the same as the input type; an arbitrary user-defined function could do all sorts of other things before/after calling `np.min`.

Even if `meta_from_array` had raised an error, this still wouldn't have worked. This logic for generating meta from the `output_dtypes` https://github.com/dask/dask/blob/ac1bd05cfd40207d68f6eb8603178d7ac0ded922/dask/array/gufunc.py#L436-L440 will rarely/never be reached, because `output_dtypes` is typically a list, not a tuple.

Overall, there is inconsistency throughout `apply_gufunc` as to whether `nout`, the type+length of `output_dtypes`, or the type+length of `metas` is the source of truth for how many items to return.


**Anything else we need to know?**:

**Environment**:

- Dask version: ac1bd05cfd40207d68f6eb8603178d7ac0ded922
- Python version: 3.8.8
- Operating System: macOS
- Install method (conda, pip, source): source



Discussion:

Resolving patch:
diff --git a/dask/array/gufunc.py b/dask/array/gufunc.py
--- a/dask/array/gufunc.py
+++ b/dask/array/gufunc.py
@@ -171,7 +171,20 @@ def _validate_normalize_axes(axes, axis, keepdims, input_coredimss, output_cored
     return input_axes, output_axes
 
 
-def apply_gufunc(func, signature, *args, **kwargs):
+def apply_gufunc(
+    func,
+    signature,
+    *args,
+    axes=None,
+    axis=None,
+    keepdims=False,
+    output_dtypes=None,
+    output_sizes=None,
+    vectorize=None,
+    allow_rechunk=False,
+    meta=None,
+    **kwargs,
+):
     """
     Apply a generalized ufunc or similar python function to arrays.
 
@@ -274,15 +287,6 @@ def apply_gufunc(func, signature, *args, **kwargs):
     .. [1] https://docs.scipy.org/doc/numpy/reference/ufuncs.html
     .. [2] https://docs.scipy.org/doc/numpy/reference/c-api/generalized-ufuncs.html
     """
-    axes = kwargs.pop("axes", None)
-    axis = kwargs.pop("axis", None)
-    keepdims = kwargs.pop("keepdims", False)
-    output_dtypes = kwargs.pop("output_dtypes", None)
-    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
     if not isinstance(signature, str):
@@ -292,38 +296,65 @@ def apply_gufunc(func, signature, *args, **kwargs):
     ## Determine nout: nout = None for functions of one direct return; nout = int for return tuples
     nout = None if not isinstance(output_coredimss, list) else len(output_coredimss)
 
-    ## Determine and handle output_dtypes
-    if output_dtypes is None:
-        if vectorize:
-            tempfunc = np.vectorize(func, signature=signature)
-        else:
-            tempfunc = func
-        output_dtypes = apply_infer_dtype(
-            tempfunc, args, kwargs, "apply_gufunc", "output_dtypes", nout
+    ## Consolidate onto `meta`
+    if meta is not None and output_dtypes is not None:
+        raise ValueError(
+            "Only one of `meta` and `output_dtypes` should be given (`meta` is preferred)."
         )
+    if meta is None:
+        if output_dtypes is None:
+            ## Infer `output_dtypes`
+            if vectorize:
+                tempfunc = np.vectorize(func, signature=signature)
+            else:
+                tempfunc = func
+            output_dtypes = apply_infer_dtype(
+                tempfunc, args, kwargs, "apply_gufunc", "output_dtypes", nout
+            )
 
-    if isinstance(output_dtypes, (tuple, list)):
+        ## Turn `output_dtypes` into `meta`
+        if (
+            nout is None
+            and isinstance(output_dtypes, (tuple, list))
+            and len(output_dtypes) == 1
+        ):
+            output_dtypes = output_dtypes[0]
+        sample = args[0] if args else None
         if nout is None:
-            if len(output_dtypes) > 1:
+            meta = meta_from_array(sample, dtype=output_dtypes)
+        else:
+            meta = tuple(meta_from_array(sample, dtype=odt) for odt in output_dtypes)
+
+    ## Normalize `meta` format
+    meta = meta_from_array(meta)
+    if isinstance(meta, list):
+        meta = tuple(meta)
+
+    ## Validate `meta`
+    if nout is None:
+        if isinstance(meta, tuple):
+            if len(meta) == 1:
+                meta = meta[0]
+            else:
                 raise ValueError(
-                    (
-                        "Must specify single dtype or list of one dtype "
-                        "for `output_dtypes` for function with one output"
-                    )
+                    "For a function with one output, must give a single item for `output_dtypes`/`meta`, "
+                    "not a tuple or list."
                 )
-            otypes = output_dtypes
-            output_dtypes = output_dtypes[0]
-        else:
-            otypes = output_dtypes
     else:
-        if nout is not None:
+        if not isinstance(meta, tuple):
+            raise ValueError(
+                f"For a function with {nout} outputs, must give a tuple or list for `output_dtypes`/`meta`, "
+                "not a single item."
+            )
+        if len(meta) != nout:
             raise ValueError(
-                "Must specify tuple of dtypes for `output_dtypes` for function with multiple outputs"
+                f"For a function with {nout} outputs, must give a tuple or list of {nout} items for "
+                f"`output_dtypes`/`meta`, not {len(meta)}."
             )
-        otypes = [output_dtypes]
 
     ## Vectorize function, if required
     if vectorize:
+        otypes = [x.dtype for x in meta] if isinstance(meta, tuple) else [meta.dtype]
         func = np.vectorize(func, signature=signature, otypes=otypes)
 
     ## Miscellaneous
@@ -420,37 +451,25 @@ 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.
-    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
-            )
+    tmp = blockwise(
+        func, loop_output_dims, *arginds, concatenate=True, meta=meta, **kwargs
+    )
 
-    if isinstance(tmp._meta, tuple):
-        metas = tmp._meta
+    # NOTE: we likely could just use `meta` instead of `tmp._meta`,
+    # but we use it and validate it anyway just to be sure nothing odd has happened.
+    metas = tmp._meta
+    if nout is None:
+        assert not isinstance(
+            metas, (list, tuple)
+        ), f"meta changed from single output to multiple output during blockwise: {meta} -> {metas}"
+        metas = (metas,)
     else:
-        metas = (tmp._meta,)
+        assert isinstance(
+            metas, (list, tuple)
+        ), f"meta changed from multiple output to single output during blockwise: {meta} -> {metas}"
+        assert (
+            len(metas) == nout
+        ), f"Number of outputs changed from {nout} to {len(metas)} during blockwise"
 
     ## Prepare output shapes
     loop_output_shape = tmp.shape
@@ -461,7 +480,6 @@ def apply_gufunc(func, signature, *args, **kwargs):
     ### *) Treat direct output
     if nout is None:
         output_coredimss = [output_coredimss]
-        output_dtypes = [output_dtypes]
 
     ## Split output
     leaf_arrs = []
diff --git a/dask/array/utils.py b/dask/array/utils.py
--- a/dask/array/utils.py
+++ b/dask/array/utils.py
@@ -65,13 +65,6 @@ def meta_from_array(x, ndim=None, dtype=None):
     if isinstance(x, type):
         x = x(shape=(0,) * (ndim or 0), dtype=dtype)
 
-    if (
-        not hasattr(x, "shape")
-        or not hasattr(x, "dtype")
-        or not isinstance(x.shape, tuple)
-    ):
-        return x
-
     if isinstance(x, list) or isinstance(x, tuple):
         ndims = [
             0
@@ -84,6 +77,13 @@ def meta_from_array(x, ndim=None, dtype=None):
         a = [a if nd == 0 else meta_from_array(a, nd) for a, nd in zip(x, ndims)]
         return a if isinstance(x, list) else tuple(x)
 
+    if (
+        not hasattr(x, "shape")
+        or not hasattr(x, "dtype")
+        or not isinstance(x.shape, tuple)
+    ):
+        return x
+
     if ndim is None:
         ndim = x.ndim
 
@@ -156,8 +156,10 @@ def compute_meta(func, _dtype, *args, **kwargs):
                 else:
                     return None
             except ValueError as e:
-                # min/max functions have no identity, attempt to use the first meta
-                if "zero-size array to reduction operation" in str(e):
+                # min/max functions have no identity, just use the same input type when there's only one
+                if len(
+                    args_meta
+                ) == 1 and "zero-size array to reduction operation" in str(e):
                     meta = args_meta[0]
                 else:
                     return None
@@ -201,7 +203,7 @@ def _not_empty(x):
 
 
 def _check_dsk(dsk):
-    """ Check that graph is well named and non-overlapping """
+    """Check that graph is well named and non-overlapping"""
     if not isinstance(dsk, HighLevelGraph):
         return
 
@@ -509,7 +511,7 @@ def asanyarray_safe(a, like, **kwargs):
 
 
 def validate_axis(axis, ndim):
-    """ Validate an input to axis= keywords """
+    """Validate an input to axis= keywords"""
     if isinstance(axis, (tuple, list)):
         return tuple(validate_axis(ax, ndim) for ax in axis)
     if not isinstance(axis, numbers.Integral):


Test patch:
diff --git a/dask/array/tests/test_array_utils.py b/dask/array/tests/test_array_utils.py
--- a/dask/array/tests/test_array_utils.py
+++ b/dask/array/tests/test_array_utils.py
@@ -36,6 +36,7 @@ def test_meta_from_array(asarray):
     assert meta_from_array(x, ndim=2).shape == (0, 0)
     assert meta_from_array(x, ndim=4).shape == (0, 0, 0, 0)
     assert meta_from_array(x, dtype="float64").dtype == "float64"
+    assert meta_from_array(x, dtype=float).dtype == "float64"
 
     x = da.ones((1,))
     assert isinstance(meta_from_array(x), np.ndarray)
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
@@ -146,14 +146,15 @@ def stats(x):
 @pytest.mark.parametrize("vectorize", [False, True])
 def test_apply_gufunc_output_dtypes_string_many_outputs(vectorize):
     def stats(x):
-        return np.mean(x, axis=-1), np.std(x, axis=-1)
+        return np.mean(x, axis=-1), np.std(x, axis=-1), np.min(x, axis=-1)
 
     a = da.random.normal(size=(10, 20, 30), chunks=(5, 5, 30))
-    mean, std = apply_gufunc(
-        stats, "(i)->(),()", a, output_dtypes=("f", "f"), vectorize=vectorize
+    mean, std, min = apply_gufunc(
+        stats, "(i)->(),(),()", a, output_dtypes=("f", "f", "f"), vectorize=vectorize
     )
     assert mean.compute().shape == (10, 20)
     assert std.compute().shape == (10, 20)
+    assert min.compute().shape == (10, 20)
 
 
 def test_apply_gufunc_pass_additional_kwargs():
@@ -161,8 +162,8 @@ def foo(x, bar):
         assert bar == 2
         return x
 
-    ret = apply_gufunc(foo, "()->()", 1.0, output_dtypes="f", bar=2)
-    assert_eq(ret, np.array(1.0, dtype="f"))
+    ret = apply_gufunc(foo, "()->()", 1.0, output_dtypes=float, bar=2)
+    assert_eq(ret, np.array(1.0, dtype=float))
 
 
 def test_apply_gufunc_02():