Для jit-класса numba как указать тип аргумента, который является функцией jit?

#python #jit #numba

Вопрос:

У меня есть коллекция функций jit, выводом которых всегда является массив numpy одного и того же типа и формы, который я хотел бы передать классу jit в качестве атрибутов, но не уверен, как определить спецификацию их типа. Похоже, в документации упоминается только использование функций jit в качестве входных данных для других функций jit.

Простой пример того, что я хочу сделать, заключается в следующем:

 @jit
def input_function(a, b):
    return a b

@jitclass(spec)
class TestClass:

    def __init__(self, arbitrary_jit_function):
        self.arbitrary_jit_function = arbitrary_jit_function
 

Как мне указать функцию arbitrary_jit_function в спецификации?