Использование функций numba в map_blocks - PullRequest
Новая
       21

Использование функций numba в map_blocks

0 голосов
/ 08 ноября 2018

Я успешно использовал map_blocks несколько раз в массивах dask. Сейчас я пытаюсь развернуть функцию numba, чтобы воздействовать на каждый блок, а также действовать и изменять один из входов.

Функция numba принимает 2 массива и обновляет второй. затем он возвращается в операторе return, чтобы в результате сделать его доступным для map_blocks.

Функция отлично работает с пустым массивом, но python просто вылетает при вызове из map_blocks. Функции numba, которые не действуют на входной массив, ведут себя нормально (хотя в этом случае трудно заставить их сделать что-нибудь полезное). Это известное ограничение? Жук? Я использую это неправильно ?! Обновление Я наконец свел это к воспроизводимому примеру с тривиальной функцией 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 получает массивы numpy с набором write = False, как вы получаете numba для выполнения какой-либо полезной работы? Вы не можете поместить строку создания массива в функцию numba, и вы не можете передать ей записываемые массивы.

Есть ли мнения о том, как этого добиться?

1 Ответ

0 голосов
/ 13 ноября 2018

Вот версия вашего кода с созданием массива, которая отлично работает в режиме numba nopython 

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=.Обратите внимание на список поддерживаемых функций , согласно документации.

...