Dask.array.套用_沿_轴：由于额外的元素([1])，使用dask.array的每一行作为另一个函数的输入失败

问题描述

我有一个形状为(62000000, 2)的大型数组(arr)，每行表示两个整数索引，我希望将它们传递给另一个函数。结构有点像这样：

def myfunc(a, b):
    return a + b
def pair_func(two_elem_arr):  # takes a 2-element vector/array
    return myfunc(*two_elem_arr)

我尝试使用dask并行化整个进程，但遇到以下问题。

当使用apply_along_axis仅打印时，生成的背景循环的第一个元素是一个神秘的[1]。当使用普通numpy：

时，完全相同的过程工作得非常好

import numpy as np
import dask.array as da

arr = np.repeat(np.arange(10), 2).reshape((10, 2))
darr = da.from_array(arr)

print('numpy version:')
np.apply_along_axis(arr=arr, axis=1, func1d=print)
print('
dask version:')
da.apply_along_axis(arr=darr, axis=1, func1d=print).compute()

输出：

numpy version:
[0 0]
[1 1]
[2 2]
[3 3]
[4 4]
[5 5]
[6 6]
[7 7]
[8 8]
[9 9]

dask version:
[1]   <------------- ?
[0 0]
[1 1]
[2 2]
[3 3]
[4 4]
[5 5]
[6 6]
[7 7]
[8 8]
[9 9]

这[1]来自哪里？

这似乎是应用所需函数失败的原因：

np.apply_along_axis(arr=arr, axis=1, func1d=pair_func)
da.apply_along_axis(arr=darr, axis=1, func1d=pair_func)

da.apply_along_axis调用导致以下异常：

TypeError: myfunc() missing 1 required positional argument: 'b'

更让我困惑的是，整个过程与da.sum(以及np.sum在同一位置)等其他函数一起工作：

da.apply_along_axis(arr=darr, axis=1, func1d=da.sum).compute()

所以，实际上有两个问题

• da.apply_along_axis(...)的输出中为什么会有[1]，如何删除
• 如果没有，是否有其他方法可以达到预期效果？

提前谢谢

---

解决方案

解决方案隐藏在documentation中。

[...]

备注

如果未提供数据类型或形状，Dask会尝试通过以下方式确定它们 在伪数组上调用函数1d。这可能会产生不正确的Dtype或Shape值， 因此，我们建议提供它们。 [...]

当我提供shape=(1,)和dtype='int'时，它工作正常。 在我看来，额外的[1]是使用伪数组确定dtype和shape的结果。不过，我认为事情不应该是这样的。

da.apply_along_axis(arr=darr, axis=1, func1d=pair_func, shape=(1,), dtype='int')