Repository: dask/dask

Issue:
da.stack dtype
Greetings dask dudes,

I'm using `dask` version `0.14.1`

I stumbled across the following behaviour with respect to the `dtype` of a **computed** `stack`, and I was curious as to whether it is as expected:

First, create some data ...

```python
import dask.array as da                           
import numpy as np                 

# integer
i = np.array(np.arange(3), dtype='i4')
# float
f = np.array(np.arange(3), dtype='f4')

di = da.from_array(i, chunks=i.shape)
df = da.from_array(f, chunks=f.shape)
```

Now, `stack` the arrays in alternate orders ...
```python
didf = da.stack([di, df], axis=0)
dfdi = da.stack([df, di], axis=0)
```
And show the results, paying attention to the lazy `dtype` and the computed `dtype` ...
```python
print(didf)
print(dfdi)
print(didf.compute().dtype)
print(dfdi.compute().dtype)
```
Yields ...
```python
dask.array<stack, shape=(2, 3), dtype=float64, chunksize=(1, 3)>
dask.array<stack, shape=(2, 3), dtype=float64, chunksize=(1, 3)>
int32
float32
```
Comparing this with `concatenate`, we see a different result, but a result I would expect ...
```python
didf = da.concatenate([di, df], axis=0)
dfdi = da.concatenate([df, di], axis=0)
print(didf)
print(dfdi)
print(didf.compute().dtype)
print(dfdi.compute().dtype)
```
Yields ...
```python
dask.array<concatenate, shape=(6,), dtype=float64, chunksize=(3,)>
dask.array<concatenate, shape=(6,), dtype=float64, chunksize=(3,)>
float64
float64
```
Is this non-symmetric `dtype` behaviour of `stack` as you would expect?

Many thanks!


Discussion:

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 9bf4edb3e..83c5160b4 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -514,7 +514,8 @@ def map_blocks(func, *args, **kwargs):
     drop_axis : number or iterable, optional
         Dimensions lost by the function.
     new_axis : number or iterable, optional
-        New dimensions created by the function.
+        New dimensions created by the function. Note that these are applied
+        after ``drop_axis`` (if present).
     name : string, optional
         The key name to use for the array. If not provided, will be determined
         by a hash of the arguments.
@@ -611,9 +612,6 @@ def map_blocks(func, *args, **kwargs):
     if isinstance(new_axis, Number):
         new_axis = [new_axis]
 
-    if drop_axis and new_axis:
-        raise ValueError("Can't specify drop_axis and new_axis together")
-
     arrs = [a for a in args if isinstance(a, Array)]
     other = [(i, a) for i, a in enumerate(args) if not isinstance(a, Array)]
 
@@ -666,15 +664,22 @@ def map_blocks(func, *args, **kwargs):
                           if i - 1 not in drop_axis), v)
                    for k, v in dsk.items())
         numblocks = [n for i, n in enumerate(numblocks) if i not in drop_axis]
-    elif new_axis:
+    if new_axis:
+        new_axis = sorted(new_axis)
+        for i in new_axis:
+            if not 0 <= i <= len(numblocks):
+                ndim = len(numblocks)
+                raise ValueError("Can't add axis %d when current "
+                                 "axis are %r. Missing axis: "
+                                 "%r" % (i, list(range(ndim)),
+                                         list(range(ndim, i))))
+            numblocks.insert(i, 1)
         dsk, old_dsk = dict(), dsk
         for key in old_dsk:
             new_key = list(key)
             for i in new_axis:
                 new_key.insert(i + 1, 0)
             dsk[tuple(new_key)] = old_dsk[key]
-        for i in sorted(new_axis):
-            numblocks.insert(i, 1)
 
     if chunks:
         if len(chunks) != len(numblocks):
@@ -702,7 +707,7 @@ def map_blocks(func, *args, **kwargs):
                                  "`chunks` kwarg.")
         if drop_axis:
             chunks2 = [c for (i, c) in enumerate(chunks2) if i not in drop_axis]
-        elif new_axis:
+        if new_axis:
             for i in sorted(new_axis):
                 chunks2.insert(i, (1,))
 
@@ -2255,6 +2260,9 @@ def stack(seq, axis=0):
     uc_args = list(concat((x, ind) for x in seq))
     _, seq = unify_chunks(*uc_args)
 
+    dt = reduce(np.promote_types, [a.dtype for a in seq])
+    seq = [x.astype(dt) for x in seq]
+
     assert len(set(a.chunks for a in seq)) == 1  # same chunks
     chunks = (seq[0].chunks[:axis] + ((1,) * n,) + seq[0].chunks[axis:])
 
@@ -2271,8 +2279,6 @@ def stack(seq, axis=0):
     dsk = dict(zip(keys, values))
     dsk2 = sharedict.merge((name, dsk), *[a.dask for a in seq])
 
-    dt = reduce(np.promote_types, [a.dtype for a in seq])
-
     return Array(dsk2, name, chunks, dtype=dt)
 
 


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index b31efda9f..8feff88cb 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -277,6 +277,15 @@ def test_stack_scalars():
     assert s.compute().tolist() == [np.arange(4).mean(), np.arange(4).sum()]
 
 
+def test_stack_promote_type():
+    i = np.arange(10, dtype='i4')
+    f = np.arange(10, dtype='f4')
+    di = da.from_array(i, chunks=5)
+    df = da.from_array(f, chunks=5)
+    res = da.stack([di, df])
+    assert_eq(res, np.stack([i, f]))
+
+
 @pytest.mark.skipif(LooseVersion(np.__version__) < '1.10.0',
                     reason="NumPy doesn't yet support stack")
 def test_stack_rechunk():
@@ -2106,9 +2115,9 @@ def test_map_blocks_with_changed_dimension():
     with pytest.raises(ValueError):
         d.map_blocks(lambda b: b.sum(axis=1), drop_axis=1, dtype=d.dtype)
 
-    # Can't use both drop_axis and new_axis
+    # Adding axis with a gap
     with pytest.raises(ValueError):
-        d.map_blocks(lambda b: b, drop_axis=1, new_axis=1)
+        d.map_blocks(lambda b: b, new_axis=(3, 4))
 
     d = da.from_array(x, chunks=(4, 8))
     e = d.map_blocks(lambda b: b.sum(axis=1), drop_axis=1, dtype=d.dtype)
@@ -2127,6 +2136,19 @@ def test_map_blocks_with_changed_dimension():
     assert e.chunks == ((1,), (4, 4), (4, 4), (1,))
     assert_eq(e, x[None, :, :, None])
 
+    # Both new_axis and drop_axis
+    d = da.from_array(x, chunks=(8, 4))
+    e = d.map_blocks(lambda b: b.sum(axis=0)[:, None, None],
+                     drop_axis=0, new_axis=(1, 2), dtype=d.dtype)
+    assert e.chunks == ((4, 4), (1,), (1,))
+    assert_eq(e, x.sum(axis=0)[:, None, None])
+
+    d = da.from_array(x, chunks=(4, 8))
+    e = d.map_blocks(lambda b: b.sum(axis=1)[:, None, None],
+                     drop_axis=1, new_axis=(1, 2), dtype=d.dtype)
+    assert e.chunks == ((4, 4), (1,), (1,))
+    assert_eq(e, x.sum(axis=1)[:, None, None])
+
 
 def test_broadcast_chunks():
     assert broadcast_chunks(((5, 5),), ((5, 5),)) == ((5, 5),)