Dask apply_along_axis 错误，与 Numpy 比较

Question

我正在尝试将函数应用于 Dask 数组 using apply_along_axis，虽然相同的函数适用于 numpy 数组，但它不适用于 Dask 数组。 这是一个例子：

import dask.array as da

w = numpy.array([[6,7,8],[9,10,11]])
q = numpy.array([[1,2,3],[4,5,6]])
s = numpy.stack([w,q])
def func(arr):
    t, y = arr[0], arr[1]
    return arr[0] + arr[1]

s_dask = da.from_array(s)

在 numpy 数组上运行 func 按预期工作，而在 Dask 数组上运行它会引发错误：IndexError: index 1 is out of bounds for axis 0 with size 1"

>>>s
array([[[ 6,  7,  8],
        [ 9, 10, 11]],

       [[ 1,  2,  3],
        [ 4,  5,  6]]])

>>>numpy.apply_along_axis(func,0,s)
array([[ 7,  9, 11],
       [13, 15, 17]])

>>>da.apply_along_axis(func,0,s_dask)
Traceback (most recent call last):
  File "<pyshell#151>", line 1, in <module>
    da.apply_along_axis(func,0,s_dask)
  File "..Python37\lib\site-packages\dask\array\routines.py", line 383, in apply_along_axis
    test_result = np.array(func1d(test_data, *args, **kwargs))
  File "<pyshell#149>", line 2, in func
    t, y = a[0],a[1]
IndexError: index 1 is out of bounds for axis 0 with size 1

我不确定我在这里做错了什么

Answer 1

Dask 数组试图找出输出数组的 dtype。为此，它通过您的函数发送一个小的空数组。那个小的空数组失败了，因为你的函数假设输入的大小至少为 2。

您可以通过显式提供 dtype 来省去 Dask 的麻烦。

da.apply_along_axis(func, 0, s_dask, dtype=s_dask.dtype)