Как указать кортеж в подписи Numba Vectorize?

Как указать кортеж в подписи Numba Vectorize?

В функции numba.vectorize вы не можете использовать кортеж. Это потому, что vectorize векторизует код для массивов этих типов.

Таким образом, использование сигнатуры float, float, tuple создает функцию, которая ожидает два массива, содержащих числа с плавающей запятой, и один массив, содержащий кортежи. Проблема в том, что для массива, содержащего кортежи, нет dtype - он может работать, если вы используете структурированный массив вместо массива, содержащего кортежи, но я этого не пробовал.

Как указать кортеж в подписи Numba jit?

Правильный способ указать UniTuple в нумба-подписи — использовать numba.types.containers.UniTuple. В твоем случае:

nb.types.containers.UniTuple(nb.types.float64, 9)

Таким образом, правильная подпись будет выглядеть примерно так:

import numba as nb

@nb.njit(
    nb.types.float64(
        nb.types.float64, 
        nb.types.float64, 
        nb.types.containers.UniTuple(nb.types.float64, 9)))
def func(f1, f2, ftuple):
    # ...
    return f1

Я часто избегаю явного ввода своих функций numba, но когда я это делаю, я нахожу очень полезным использовать numba.typeof, например:

>>> nb.typeof((1.0, ) * 9)
tuple(float64 x 9)

>>> type(nb.typeof((1.0, ) * 9))
numba.types.containers.UniTuple

>>> help(type(nb.typeof((1.0, ) * 9)))  # I shortened the result:
Help on class UniTuple in module numba.types.containers:

class UniTuple(BaseAnonymousTuple, _HomogeneousTuple, numba.types.abstract.Sequence)
 |  UniTuple(*args, **kwargs)
 |  
 |  Type class for homogeneous tuples.
 |  
 |  Methods defined here:
 |  
 |  __init__(self, dtype, count)
 |      Initialize self.  See help(type(self)) for accurate signature.

Итак, вся информация есть: это numba.types.containes.UniTuple, и вы создаете его экземпляр с двумя аргументами: dtype (здесь float64) и числом (в данном случае 9).

Если вы хотите векторизовать только массивы с плавающей запятой

Если вы не хотите векторизовать функцию для аргумента кортежа, вы можете просто создать векторизованную функцию внутри другой функции и вызвать ее там:

import numba as nb
import numpy as np

def func(E, L, fparams):
    @nb.vectorize(['float64(float64, float64)'])
    def fn_vec(e, l):
        return e + l + fparams[1]  # just to illustrate that the tuple is available
    return fn_vec(E, L)

Это делает кортеж доступным внутри функции vectorized. Однако он должен создавать внутреннюю функцию и компилировать ее каждый раз, когда вы вызываете внешнюю функцию, поэтому на самом деле это может быть медленнее. Я также не уверен, что это будет работать с target="cuda", возможно, вам придется проверить это самостоятельно.