Numba字典：JIT（）装饰器中的签名

Question

我的函数将一个numpy数组列表和一个字典（或词典列表）作为输入参数，并返回值列表。 numpy数组的列表很长，并且数组的形状可能不同。尽管我可以分别传递numpy数组，但出于内部管理的目的，我真的很想形成一个numpy数组的元组并将其照常传递给我的函数。没有字典（根据numba> = 0.43特殊构造的字典），整个设置工作正常-请参见下面的脚本。因为输入和输出的结构是元组形式的，所以JIT需要签名-没有它就无法弄清楚数据结构的类型。但是，无论我如何尝试在JIT装饰器中声明字典'd'，我都无法使脚本正常工作。如果有想法或解决方案，请提供帮助。

非常感谢

'''蟒：

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键一起使用...形式示例如下：result = (summ * d['k1'], prod * d['k2'])

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

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

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(5)+11.0
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


# ~~ Run function and print results
print(somefunction(arg, d))

Answer 1

我正在使用版本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)

供参考，请从官方文档中查看this example。

更新：在您的情况下，您可以简单地让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))

您可以查看官方文档here：

除非必要，建议让Numba通过使用@jit的无签名变体来推断参数类型。

我尝试了各种方法，但这是解决您指定问题的唯一方法。