如何在 numba 中使用自定义类型为方法指定函数签名

Question

我使用 numba.jitclass 装饰器来标记我的类以进行优化。

我不知道如何指定要优化的run 方法的签名。该方法将ConvertedDocument 对象数组作为参数。似乎 numba 无法自行确定数组类型，因为当我尝试在 nopython 模式下调用 run 方法时出现以下错误：

Traceback (most recent call last):
  File "numba_test.py", line 53, in <module>
    print run(a)
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/dispatcher.py", line 310, in _compile_for_args
    raise e
numba.errors.TypingError: Caused By:
Traceback (most recent call last):
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/compiler.py", line 230, in run
    stage()
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/compiler.py", line 444, in stage_nopython_frontend
    self.locals)
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/compiler.py", line 800, in type_inference_stage
    infer.propagate()
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/typeinfer.py", line 767, in propagate
    raise errors[0]
TypingError: Internal error at <numba.typeinfer.ExhaustIterConstraint object at 0x788cc9572d50>:
--%<-----------------------------------------------------------------
Traceback (most recent call last):
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/typeinfer.py", line 128, in propagate
    constraint(typeinfer)
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/typeinfer.py", line 264, in __call__
    raise TypingError("failed to unpack {}".format(tp), loc=self.loc)
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/contextlib.py", line 35, in __exit__
    self.gen.throw(type, value, traceback)
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/errors.py", line 249, in new_error_context
    six.reraise(type(newerr), newerr, sys.exc_info()[2])
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/errors.py", line 243, in new_error_context
    yield
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/typeinfer.py", line 264, in __call__
    raise TypingError("failed to unpack {}".format(tp), loc=self.loc)
InternalError: local variable 'tp' referenced before assignment
[1] During: typing of exhaust iter at numba_test.py (40)
--%<-----------------------------------------------------------------

File "numba_test.py", line 40

Failed at nopython (nopython frontend)
Internal error at <numba.typeinfer.ExhaustIterConstraint object at 0x788cc9572d50>:
--%<-----------------------------------------------------------------
Traceback (most recent call last):
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/typeinfer.py", line 128, in propagate
    constraint(typeinfer)
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/typeinfer.py", line 264, in __call__
    raise TypingError("failed to unpack {}".format(tp), loc=self.loc)
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/contextlib.py", line 35, in __exit__
    self.gen.throw(type, value, traceback)
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/errors.py", line 249, in new_error_context
    six.reraise(type(newerr), newerr, sys.exc_info()[2])
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/errors.py", line 243, in new_error_context
    yield
  File "/home/clasocki/anaconda2/envs/my_numba_env/lib/python2.7/site-packages/numba/typeinfer.py", line 264, in __call__
    raise TypingError("failed to unpack {}".format(tp), loc=self.loc)
InternalError: local variable 'tp' referenced before assignment
[1] During: typing of exhaust iter at numba_test.py (40)
--%<-----------------------------------------------------------------

File "numba_test.py", line 40

This error may have been caused by the following argument(s):
- argument 0: Unsupported array dtype: object

我是这样指定 numba 装饰器的：

spec = [
    ('profile', numba.typeof(numpy.asarray([1.0, 2.0]))),
    ('word_weights', numba.typeof(numpy.asarray([(1.0,2.0)])))
]

@numba.jitclass(spec)
class ConvertedDocument(object):
    def __init__(self, profile, word_weights):
        self.profile = profile
        self.word_weights = word_weights

@numba.jit(nopython=True,cache=True)
def run(docs):
    s = 0
    for doc in docs: #array of documents
        for w_id, weight in doc.word_weights: #accessing document's property
            s += weight

    return s

这就是run方法的调用方式：

x = numpy.asarray([1.0, 2.0])
y = numpy.asarray([(1.0,2.0), (3.0,4.0)])
a = numpy.asarray([ConvertedDocument(x,y)])
print run(a)

如果anumpy数组被替换成Python列表，异常如下：

Failed at nopython (nopython mode backend)
reflected list(instance.jitclass.ConvertedDocument#3bffb70<profile:array(float64, 1d, C),word_weights:array(float64, 2d, C)>): unsupported nested memory-managed object

有谁知道在使用自定义类型时如何指定方法签名或者是否支持对对象数组的迭代？

Answer 1

问题似乎是您不能在 jitclass 对象上调用 np.nditer，这是有道理的，因为 jitclass 不可迭代。它将数据存储为数组结构（和其他数据类型），而不是结构数组。您正试图将其用作后者。如果除了你的两个数组属性之外你还有一堆标量数据属性或不同大小的数组，你将如何迭代 jitclass 对象将是模棱两可的。

错误信息确实不清楚。我的建议是迭代您直接需要的word_weights 的索引。