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

2022-04-14 00:00:00 python numpy dask

问题描述

我有一个形状为(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]是使用伪数组确定dtypeshape的结果。不过,我认为事情不应该是这样的。

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

相关文章