dask数组分块如何在三个或更多个维度上起作用？

Question

考虑以下代码片段，其中我试图创建一个大小为(3, 4, 5)的不对称但重复的块的dask数组：

import numpy as np
import dask.array as da
a = np.random.randint(0, 9, (3, 12, 10))
d = da.from_array(a, chunks=(3, 4, 5))

以上代码段不会引发任何错误/警告。但是当我尝试执行以下操作时：

r = d.map_blocks(np.sum)
out = r.compute()

它引发以下错误：

python3.7/site-packages/dask/array/core.py in <listcomp>(.0)
   4099 
   4100     while isinstance(arrays, (list, tuple)):
-> 4101         result.append(tuple([shape(deepfirst(a))[dim] for a in arrays]))
   4102         arrays = arrays[0]
   4103         dim += 1

IndexError: tuple index out of range

我在做什么错？

Answer 1

如果您检查map_block文档here。在chunks参数描述下，它表示：

“如果未提供，则假定结果数组具有与第一个输入数组相同的块结构。”

所以您的sum函数应该是这样的：

def compute_block_sum(block):
    return np.array([[np.sum(block)]])[:,None]

然后您做

d.map_blocks(compute_block_sum).compute()