在map_blocks中使用numba函数



我已经在dask数组上成功地使用了几次map_blocks。我现在正试图部署一个numba函数来处理每个块,并处理和更改其中一个输入。

numba函数接收2个numpy数组,并更新第二个数组。然后在return语句中返回该值,从而使其可用于mapblocks。

该函数在numpy数组上运行良好,但python只是在从map_blocks调用它时崩溃。不作用于输入数组的numba函数行为正常(尽管在这种情况下很难让它们做任何有用的事情(。这是已知的限制吗?虫子?我用错了吗?!Update我终于用一个微不足道的numba函数把它归结为一个可复制的例子,我对这个问题有了更清晰的了解。然而,我仍然不清楚如何解决这个问题。这是代码:

import numpy as np
from numba import jit, float64, int64
from dask.distributed import Client, LocalCluster
import dask.array as da
cluster=LocalCluster()
c=Client(cluster)
size=int(1e5)
a=np.arange(size,dtype='float64')
b=np.zeros((size,),dtype='float64')
dista=da.from_array(a,chunks=size//4)
distb=da.from_array(b,chunks=size//4)
@jit(float64[:](float64[:],float64[:]))
def crasher(x,y):
for i in range(x.shape[0]):
y[i]=x[i]*2
return y
distc=da.map_blocks(crasher,dista,distb,dtype='float64')
c=distc.compute() #it all crashes at this point

我现在得到了一个更容易理解的错误,而不仅仅是一个直接的崩溃:

TypeError: No matching definition for argument type(s) readonly array(float64, 1d, C), readonly array(float64, 1d, C)

所以,如果numba正在接收write=False设置的numpy数组,您如何让numba做任何有用的工作?不能在numba函数中放入数组创建行,也不能为其提供可写数组。

对如何实现这一目标有什么看法?

这是一个创建数组的代码版本,它在numba nopyson模式下运行良好

import numpy as np
from numba import jit, float64, int64
from dask.distributed import Client, LocalCluster
import dask.array as da
cluster=LocalCluster()
c=Client(cluster)
size=int(1e5)
a=np.arange(size,dtype='float64')
dista=da.from_array(a,chunks=size//4)
@jit(nopython=True)
def crasher(x):
y = np.empty_like(x)
for i in range(x.shape[0]):
y[i]=x[i]*2
return y
distc=da.map_blocks(crasher,dista,dtype='float64')
c=distc.compute()

注意y=行。根据文档,请注意支持的numpy函数列表。

最新更新