Словарь Numba: подпись в JIT () декораторе - PullRequest
Новая
       18

Словарь Numba: подпись в JIT () декораторе

3 голосов
/ 28 сентября 2019

Моя функция принимает список числовых массивов и словарь (или список словарей) в качестве входных аргументов и возвращает список значений.Список массивов numpy длинный, и массивы могут иметь различную форму.Хотя я могу передавать отдельные массивы отдельно, для служебных целей я действительно хотел бы сформировать кортеж массивов и передать их как таковые в свою функцию.Без словаря (который специально сформирован в соответствии с numba> = 0,43) вся установка работает нормально - см. Скрипт ниже.Поскольку структура ввода и вывода имеет форму кортежа, JIT требует подписи - он не может определить тип структуры данных без нее.Однако независимо от того, как я пытаюсь объявить свой словарь 'd' в JIT-декораторе, мне не удается заставить скрипт работать.Пожалуйста, помогите с идеями или решением, если оно существует.

Большое спасибо

'' 'python: 

import numpy as np
from numba import njit
from numba import types
from numba.typed import Dict

@njit(  'Tuple( (f8,f8) )(Tuple( (f8[:],f8[:]) ))'  )

def somefunction(lst_arr):
    arr1, arr2 = lst_arr

    summ = 0
    prod = 1
    for i in arr2:
        summ += i
    for j in arr1:
        prod *= j

    result = (summ,prod)
    return result

a = np.arange(5)+1.0
b = np.arange(5)+11.0
arg = (a,b)
print(a,b)

print(somefunction(arg))


# ~~ The Dict.empty() constructs a typed dictionary.
d = Dict.empty(
    key_type=types.unicode_type,
    value_type=types.float64,)

d['k1'] = 1.5
d['k2'] = 0.5

' ''

Я ожидаюпередать d-словарь в somefunction и использовать его внутри с ключами dict ... Образец приведен ниже:

1 Ответ

1 голос
/ 29 сентября 2019

Я использую версию 0.45.1.Вы можете просто передать словарь, не объявляя тип в словаре: 

d = Dict.empty(
    key_type=types.unicode_type,
    value_type=types.float64[:],
)
d['k1'] = np.arange(5) + 1.0
d['k2'] = np.arange(5) + 11.0

# Numba will infer the type on it's own.
@njit
def somefunction2(d):
    prod = 1

    # I am assuming you want sum of second array and product of second
    result = (d['k2'].sum(), d['k1'].prod())

    return result

print(somefunction(d))
# Output : (65.0, 120.0)

Для справки вы проверяете этот пример из официальной документации.

Обновление :
В вашем случае вы можете просто позволить jit выводить типы самостоятельно, и это должно работать, у меня работает следующий код: 

import numpy as np
from numba import njit
from numba import types
from numba.typed import Dict
from numba.types import DictType

# Let jit infer the types on it's own
@njit
def somefunction(lst_arr, mydict):
    arr1, arr2 = lst_arr
    summ = 0
    prod = 1
    for i in arr2:
        summ += i
    for j in arr1:
        prod *= j

    result = (summ*mydict['k1'],prod*mydict['k2'])
    return result

# ~~ Input numpy arrays
a = np.arange(5)+1.0
b = np.arange(10)+11.0  #<--------------- This is of different shape 
arg = (a,b)

# ~~ Input dictionary for the function 
d = Dict.empty(
    key_type=types.unicode_type,
    value_type=types.float64)

d['k1'] = 1.5
d['k2'] = 0.5


# This works now
print(somefunction(arg, d))

Вы можете видетьофициальная документация здесь :

Если в этом нет необходимости, рекомендуется разрешить Numba выводить типы аргументов, используя вариант без подписи @ jit.

Я пробовал разные методы, но это единственный, который работал для указанной вами проблемы.

...