Repository: dask/dask

Issue:
stack followed by transpose fails
``` python
import numpy as np
import dask
import dask.array as da

da.set_options(get=dask.get)
x = da.from_array(np.array([1]), blockshape=(1,))
da.stack([x]).transpose().compute()
```

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-14-2da89db80151> in <module>()
      5 da.set_options(get=dask.get)
      6 x = da.from_array(np.array([1]), blockshape=(1,))
----> 7 da.stack([x]).transpose().compute()

/Users/shoyer/dev/dask/dask/array/core.py in compute(self, **kwargs)
    595     @wraps(compute)
    596     def compute(self, **kwargs):
--> 597         return compute(self, **kwargs)
    598 
    599     __float__ = __int__ = __bool__ = __complex__ = compute

/Users/shoyer/dev/dask/dask/array/core.py in compute(*args, **kwargs)
    430     dsk = merge(*[arg.dask for arg in args])
    431     keys = [arg._keys() for arg in args]
--> 432     results = get(dsk, keys, **kwargs)
    433 
    434     results2 = [rec_concatenate(x) if arg.shape else unpack_singleton(x)

/Users/shoyer/dev/dask/dask/array/core.py in get(dsk, keys, get, **kwargs)
    823     dsk3 = remove_full_slices(dsk2)
    824     dsk4 = inline_functions(dsk3, fast_functions=fast_functions)
--> 825     return get(dsk4, keys, **kwargs)
    826 
    827 

/Users/shoyer/dev/dask/dask/core.pyc in get(d, key, get, concrete, **kwargs)
    111         v = (get(d, k, get=get, concrete=concrete) for k in key)
    112         if concrete:
--> 113             v = list(v)
    114     elif ishashable(key) and key in d:
    115         v = d[key]

/Users/shoyer/dev/dask/dask/core.pyc in <genexpr>((k,))
    109     get = get or _get
    110     if isinstance(key, list):
--> 111         v = (get(d, k, get=get, concrete=concrete) for k in key)
    112         if concrete:
    113             v = list(v)

/Users/shoyer/dev/dask/dask/core.pyc in get(d, key, get, concrete, **kwargs)
    111         v = (get(d, k, get=get, concrete=concrete) for k in key)
    112         if concrete:
--> 113             v = list(v)
    114     elif ishashable(key) and key in d:
    115         v = d[key]

/Users/shoyer/dev/dask/dask/core.pyc in <genexpr>((k,))
    109     get = get or _get
    110     if isinstance(key, list):
--> 111         v = (get(d, k, get=get, concrete=concrete) for k in key)
    112         if concrete:
    113             v = list(v)

/Users/shoyer/dev/dask/dask/core.pyc in get(d, key, get, concrete, **kwargs)
    111         v = (get(d, k, get=get, concrete=concrete) for k in key)
    112         if concrete:
--> 113             v = list(v)
    114     elif ishashable(key) and key in d:
    115         v = d[key]

/Users/shoyer/dev/dask/dask/core.pyc in <genexpr>((k,))
    109     get = get or _get
    110     if isinstance(key, list):
--> 111         v = (get(d, k, get=get, concrete=concrete) for k in key)
    112         if concrete:
    113             v = list(v)

/Users/shoyer/dev/dask/dask/core.pyc in get(d, key, get, concrete, **kwargs)
    122         if get is _get:
    123             # use non-recursive method by default
--> 124             return _get_task(d, v)
    125         func, args = v[0], v[1:]
    126         args2 = [get(d, arg, get=get, concrete=False) for arg in args]

/Users/shoyer/dev/dask/dask/core.pyc in _get_task(d, task, maxdepth)
     59         func, args, results = stack[-1]
     60         if not args:
---> 61             val = func(*results)
     62             if len(stack) == 1:
     63                 return val

/Users/shoyer/miniconda/envs/xray-dev/lib/python2.7/site-packages/toolz/functoolz.pyc in __call__(self, *args, **kwargs)
    214     def __call__(self, *args, **kwargs):
    215         try:
--> 216             return self._partial(*args, **kwargs)
    217         except TypeError:
    218             # If there was a genuine TypeError

/Users/shoyer/miniconda/envs/xray-dev/lib/python2.7/site-packages/numpy/core/fromnumeric.pyc in transpose(a, axes)
    535     except AttributeError:
    536         return _wrapit(a, 'transpose', axes)
--> 537     return transpose(axes)
    538 
    539 

ValueError: axes don't match array
```



Discussion:
OK, definitely an issue.  It appears that it thinks the shape of `da.stack([x]) == (1, 1)`

Oh, perhaps it should be.  The individual blocks presumably need to be expanded out along the new dimension.


Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -890,8 +890,12 @@ def stack(seq, axis=0):
     keys = list(product([name], *[range(len(bd)) for bd in blockdims]))
 
     names = [a.name for a in seq]
-    values = [(names[key[axis+1]],) + key[1:axis + 1] + key[axis + 2:]
+    inputs = [(names[key[axis+1]],) + key[1:axis + 1] + key[axis + 2:]
                 for key in keys]
+    values = [(getitem, inp, (slice(None, None, None),) * axis
+                           + (None,)
+                           + (slice(None, None, None),) * (ndim - axis))
+                for inp in inputs]
 
     dsk = dict(zip(keys, values))
     dsk2 = merge(dsk, *[a.dask for a in seq])


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -205,22 +205,30 @@ def test_stack():
 
     s = stack([a, b, c], axis=0)
 
+    colon = slice(None, None, None)
+
     assert s.shape == (3, 4, 6)
     assert s.blockdims == ((1, 1, 1), (2, 2), (3, 3))
-    assert s.dask[(s.name, 0, 1, 0)] == ('A', 1, 0)
-    assert s.dask[(s.name, 2, 1, 0)] == ('C', 1, 0)
+    assert s.dask[(s.name, 0, 1, 0)] == (getitem, ('A', 1, 0),
+                                          (None, colon, colon))
+    assert s.dask[(s.name, 2, 1, 0)] == (getitem, ('C', 1, 0),
+                                          (None, colon, colon))
 
     s2 = stack([a, b, c], axis=1)
     assert s2.shape == (4, 3, 6)
     assert s2.blockdims == ((2, 2), (1, 1, 1), (3, 3))
-    assert s2.dask[(s2.name, 0, 1, 0)] == ('B', 0, 0)
-    assert s2.dask[(s2.name, 1, 1, 0)] == ('B', 1, 0)
+    assert s2.dask[(s2.name, 0, 1, 0)] == (getitem, ('B', 0, 0),
+                                            (colon, None, colon))
+    assert s2.dask[(s2.name, 1, 1, 0)] == (getitem, ('B', 1, 0),
+                                            (colon, None, colon))
 
     s2 = stack([a, b, c], axis=2)
     assert s2.shape == (4, 6, 3)
     assert s2.blockdims == ((2, 2), (3, 3), (1, 1, 1))
-    assert s2.dask[(s2.name, 0, 1, 0)] == ('A', 0, 1)
-    assert s2.dask[(s2.name, 1, 1, 2)] == ('C', 1, 1)
+    assert s2.dask[(s2.name, 0, 1, 0)] == (getitem, ('A', 0, 1),
+                                            (colon, colon, None))
+    assert s2.dask[(s2.name, 1, 1, 2)] == (getitem, ('C', 1, 1),
+                                            (colon, colon, None))
 
     assert raises(ValueError, lambda: stack([a, b, c], axis=3))
 
@@ -230,6 +238,14 @@ def test_stack():
             stack([a, b, c], axis=2).blockdims
 
 
+def test_short_stack():
+    x = np.array([1])
+    d = da.from_array(x, blockshape=(1,))
+    s = da.stack([d])
+    assert s.shape == (1, 1)
+    assert get(s.dask, s._keys())[0][0].shape == (1, 1)
+
+
 def test_concatenate():
     a, b, c = [Array(getem(name, blockshape=(2, 3), shape=(4, 6)),
                      name, shape=(4, 6), blockshape=(2, 3))