apply_ufunc gives wrong dtype with dask=parallelized and vectorized=True

[ulijh]

Applying a function to a data array with dtype = complex returns one with dtype = float. It seems to work before commit 17b70ca.

MCVE Code Sample

import numpy as np
import xarray as xr

def func(x):
    return np.sum(x ** 2)

da = xr.DataArray(np.arange(2*3*4).reshape(2,3,4))
da = da + 1j * da
da = da.chunk(dict(dim_1=1))

da2 = xr.apply_ufunc(
    func,
    da,
    vectorize=True,
    dask="parallelized",
    output_dtypes=[da.dtype],
)

assert da2.dtype == da.dtype, "wrong dtype"

Expected Output

da and da2 should both have the same dtype=complex.

Problem Description

To me it seems to me that the kwarg meta=None somehow causes dask to allocate a float array and ignore the dtype kwargs (which seems to be carried through correctly down to dask.array.blockwise.blockwise()) . I'm not familiar with the apply_ufing and the dask code, so I can't tell on which end the bug sits.

Versions

Output of xr.show_versions()

INSTALLED VERSIONS

commit: None
python: 3.8.2 (default, Apr 8 2020, 14:31:25)
[GCC 9.3.0]
python-bits: 64
OS: Linux
OS-release: 5.6.5-arch3-1
machine: x86_64
processor:
byteorder: little
LC_ALL: None
LANG: de_DE.utf8
LOCALE: de_DE.UTF-8
libhdf5: 1.10.5
libnetcdf: 4.7.3

xarray: 0.15.2.dev47+g33a66d63
pandas: 1.0.3
numpy: 1.18.3
scipy: 1.4.1
netCDF4: 1.5.3
pydap: None
h5netcdf: 0.7.4
h5py: 2.10.0
Nio: None
zarr: None
cftime: 1.1.1.2
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: 1.3.2
dask: 2.12.0
distributed: 2.14.0
matplotlib: 3.2.1
cartopy: 0.17.0
seaborn: 0.10.0
numbagg: None
pint: None
setuptools: 46.1.3
pip: 20.0.2
conda: None
pytest: 5.4.1
IPython: 7.13.0
sphinx: 3.0.2

[mathause]

It works when you set vectorize=False

da2 = xr.apply_ufunc( 
    func, 
    da, 
    vectorize=False, 
    dask="parallelized", 
    output_dtypes=[da.dtype], 
     
)                                                                                                                                                                                       
assert da2.dtype == da.dtype, "wrong dtype"        

or when you pass your own meta:

da2 = xr.apply_ufunc( 
    func, 
    da, 
    vectorize=True, 
    dask="parallelized", 
    output_dtypes=[da.dtype], 
    meta=da 
) 

assert da2.dtype == da.dtype, "wrong dtype"

This also goes wrong if the DataArray has another dtype, e.g. int:

da = xr.DataArray(np.arange(2*3*4).reshape(2,3,4)
da = da.chunk(dict(dim_1=1))
da2 = xr.apply_ufunc(
    func,
    da,
    vectorize=True,
    dask="parallelized",
    output_dtypes=[da.dtype],
)

assert da2.dtype == da.dtype, "wrong dtype"

Indeed, the dtype of meta takes precedence over the dtype (https://github.com/dask/dask/blob/25005e19cc30e8b2877d4dadbaef378ee912bdc0/dask/array/core.py#L1022):

meta : empty ndarray
    empty ndarray created with same NumPy backend, ndim and dtype as the
    Dask Array being created (overrides dtype)

[dcherian]

@mathuse is right.

The solution is to use dtype when we create meta for vectorized functions here:

xarray/xarray/core/computation.py

Lines 1008 to 1011 in 3820fb7

	if meta is None:
	# set meta=np.ndarray by default for numpy vectorized functions
	# work around dask bug computing meta with vectorized functions: GH5642
	meta = np.ndarray

@ulijh or @mathause , are either of you up for sending in a PR?

For now the workaround is to pass your own meta = np.ndarray((0,0), dtype=da.dtype)