Using numba functions in map_blocks - dask

Question

I have successfully used map_blocks a few times on dask arrays. I'm now trying to deploy a numba function to act on each block, and to act and change one of the inputs.
The numba function takes in 2 numpy arrays, and updates the second one. this is then returned in the return statement to make it available to map_blocks as a result.
The function works fine on a numpy array, but python just crashes when calling it from map_blocks. numba functions that do not act on an input array behave normally (although it is difficult to get them to do anything useful in this case).
Is this a known limitation? A bug? Am I using it wrong?!
Update
I've finally boiled it down to a reproducible example with a trivial numba function, and I get a clearer idea of the problem. However I'm still unclear on how to resolve the issue. Here's the code:
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
And I now get a more comprehensible error rather than just a straight up crash:
TypeError: No matching definition for argument type(s) readonly array(float64, 1d, C), readonly array(float64, 1d, C)
So if numba is receiving numpy arrays with write=False set, how do you get numba to do any useful work? You can't put an array creation line in the numba function, and you can't feed it writeable arrays.
Any views on how to achieve this?

Answer 1

Here is a version of your code with array creation, which runs fine with numba nopython mode
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()
Note the y= line. Note the list of numpy functions supported, according to the documentation.