Fix meta_from_array to support Xarray test suite (dask#4938)

mrocklin · web-flow · commit c437e63e6356 · 2019-06-16T10:42:16.000+02:00
Fixes pydata/xarray#3009
diff --git a/dask/array/core.py b/dask/array/core.py
@@ -888,8 +888,7 @@ def __new__(cls, dask, name, chunks, dtype=None, meta=None, shape=None):
         if dtype:
             self._meta = np.empty((0,) * self.ndim, dtype=dtype)
         else:
-            from .utils import meta_from_array
-            self._meta = meta_from_array(meta, meta.ndim)
+            self._meta = meta_from_array(meta)
 
         for plugin in config.get('array_plugins', ()):
             result = plugin(self)
@@ -2984,8 +2983,7 @@ def concatenate(seq, axis=0, allow_unknown_chunksizes=False):
     if not seq:
         raise ValueError("Need array(s) to concatenate")
 
-    from .utils import meta_from_array
-    meta = np.concatenate([meta_from_array(s, s.ndim) for s in seq])
+    meta = np.concatenate([meta_from_array(s) for s in seq])
 
     # Promote types to match meta
     seq = [a.astype(meta.dtype) for a in seq]
@@ -4167,3 +4165,6 @@ def from_npy_stack(dirname, mmap_mode='r'):
     dsk = dict(zip(keys, values))
 
     return Array(dsk, name, chunks, dtype)
+
+
+from .utils import meta_from_array
diff --git a/dask/array/gufunc.py b/dask/array/gufunc.py
@@ -10,7 +10,7 @@
     from toolz import concat, merge, unique
 
 from .core import Array, asarray, blockwise, getitem, apply_infer_dtype
-from .utils import normalize_meta
+from .utils import meta_from_array
 from ..highlevelgraph import HighLevelGraph
 from ..core import flatten
 
@@ -398,7 +398,7 @@ def apply_gufunc(func, signature, *args, **kwargs):
         leaf_name = "%s_%d-%s" % (name, i, token)
         leaf_dsk = {(leaf_name,) + key[1:] + core_chunkinds: ((getitem, key, i) if nout else key) for key in keys}
         graph = HighLevelGraph.from_collections(leaf_name, leaf_dsk, dependencies=[tmp])
-        meta = normalize_meta(tmp._meta, len(output_shape), dtype=odt)
+        meta = meta_from_array(tmp._meta, len(output_shape), dtype=odt)
         leaf_arr = Array(graph,
                          leaf_name,
                          chunks=output_chunks,
diff --git a/dask/array/linalg.py b/dask/array/linalg.py
@@ -798,9 +798,9 @@ def lu(a):
                 # l_permuted is not referred in upper triangulars
 
     pp, ll, uu = scipy.linalg.lu(np.ones(shape=(1, 1), dtype=a.dtype))
-    pp_meta = meta_from_array(a, a.ndim, dtype=pp.dtype)
-    ll_meta = meta_from_array(a, a.ndim, dtype=ll.dtype)
-    uu_meta = meta_from_array(a, a.ndim, dtype=uu.dtype)
+    pp_meta = meta_from_array(a, dtype=pp.dtype)
+    ll_meta = meta_from_array(a, dtype=ll.dtype)
+    uu_meta = meta_from_array(a, dtype=uu.dtype)
 
     graph = HighLevelGraph.from_collections(name_p, dsk, dependencies=[a])
     p = Array(graph, name_p, shape=a.shape, chunks=a.chunks, meta=pp_meta)
@@ -1048,7 +1048,7 @@ def _cholesky(a):
     graph_upper = HighLevelGraph.from_collections(name_upper, dsk, dependencies=[a])
     graph_lower = HighLevelGraph.from_collections(name, dsk, dependencies=[a])
     cho = scipy.linalg.cholesky(np.array([[1, 2], [2, 5]], dtype=a.dtype))
-    meta = meta_from_array(a, a.ndim, dtype=cho.dtype)
+    meta = meta_from_array(a, dtype=cho.dtype)
 
     lower = Array(graph_lower, name, shape=a.shape, chunks=a.chunks, meta=meta)
     # do not use .T, because part of transposed blocks are already calculated
diff --git a/dask/array/tests/test_array_utils.py b/dask/array/tests/test_array_utils.py
@@ -0,0 +1,32 @@
+import numpy as np
+import pytest
+
+from dask.array.utils import meta_from_array
+
+asarrays = [np.asarray]
+
+try:
+    import sparse
+    asarrays.append(sparse.COO.from_numpy)
+except ImportError:
+    pass
+
+try:
+    import cupy
+    asarrays.append(cupy.asarray)
+except ImportError:
+    pass
+
+
+@pytest.mark.parametrize("asarray", asarrays)
+def test_meta_from_array(asarray):
+    x = np.ones((1, 2, 3), dtype='float32')
+    x = asarray(x)
+
+    assert meta_from_array(x).shape == (0, 0, 0)
+    assert meta_from_array(x).dtype == 'float32'
+    assert type(meta_from_array(x)) is type(x)
+
+    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"
diff --git a/dask/array/utils.py b/dask/array/utils.py
@@ -28,22 +28,19 @@ def normalize_to_array(x):
         return x
 
 
-def normalize_meta(x, ndim, dtype=None):
-    if ndim > x.ndim:
-        meta = x[(Ellipsis, ) + tuple(None for _ in range(ndim - x.ndim))]
-        meta = meta[tuple(slice(0, 0, None) for _ in range(meta.ndim))]
-    elif ndim < x.ndim:
-        meta = np.sum(x, axis=tuple(d for d in range((x.ndim - ndim))))
-    else:
-        meta = x
-
-    if dtype:
-        meta = meta.astype(dtype)
-
-    return meta
-
-
-def meta_from_array(x, ndim, dtype=None):
+def meta_from_array(x, ndim=None, dtype=None):
+    """ Normalize an array to appropriate meta object
+
+    Parameters
+    ----------
+    x: array-like
+    ndim: int
+    dtype: dtype
+
+    Returns
+    -------
+    array-like
+    """
     if isinstance(x, list) or isinstance(x, tuple):
         ndims = [0 if isinstance(a, numbers.Number)
                  else a.ndim if hasattr(a, 'ndim') else len(a) for a in x]
@@ -53,11 +50,28 @@ def meta_from_array(x, ndim, dtype=None):
     # x._meta must be a Dask Array, some libraries (e.g. zarr) implement a
     # _meta attribute that are incompatible with Dask Array._meta
     if hasattr(x, '_meta') and isinstance(x, Array):
-        meta = x._meta
-    else:
+        x = x._meta
+
+    if ndim is None:
+        ndim = x.ndim
+
+    try:
         meta = x[tuple(slice(0, 0, None) for _ in range(x.ndim))]
+        if meta.ndim != ndim:
+            if ndim > x.ndim:
+                meta = meta[(Ellipsis, ) + tuple(None for _ in range(ndim - meta.ndim))]
+                meta = meta[tuple(slice(0, 0, None) for _ in range(meta.ndim))]
+            elif ndim == 0:
+                meta = meta.sum()
+            else:
+                meta = meta.reshape((0,) * ndim)
+    except Exception:
+        meta = np.empty((0,) * ndim, dtype=dtype or x.dtype)
+
+    if dtype and meta.dtype != dtype:
+        meta = meta.astype(dtype)
 
-    return normalize_meta(meta, ndim, dtype)
+    return meta
 
 
 def allclose(a, b, equal_nan=False, **kwargs):
diff --git a/dask/array/wrap.py b/dask/array/wrap.py
@@ -73,7 +73,7 @@ def wrap_func_like(func, *args, **kwargs):
     Transform np creation function into blocked version
     """
     x = args[0]
-    meta = meta_from_array(x, x.ndim)
+    meta = meta_from_array(x)
     shape = kwargs.get('shape', x.shape)
 
     parsed = _parse_wrap_args(func, args, kwargs, shape)
