Как указать кортеж в сигнатуре Numba Vectorize? - PullRequest
Новая
       27

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

2 голосов
/ 19 апреля 2019

Я определяю функцию и хочу использовать Numba Vectorize, чтобы ускорить ее с помощью cuda.У меня проблемы с сигнатурой функции.Функция вернет значение float64.Я хочу передать два значения float64, которые будут векторизованы, и дополнительно 9-кратный набор значений float64, которые будут скалярами.

Вот мой заголовок функции: 

from numba import vectorize

@vectorize(['float64(float64, float64, UniTuple(float64, 9))'], target='cuda')
def fn_vec(E, L, fparams):
    # calculations... 
    return result

, но это выдает ошибку: 

TypeError: data type "(float64 x 9)" not understood

Я пробовал много вариантов, включая (float64, ..., float64) вместо UniTuple (), но не может заставить что-либо работать.Как мне это сделать?

1 Ответ

3 голосов
/ 20 апреля 2019

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

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

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

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

Правильный способ указать UniTuple в подписи Numba с помощью 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)

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

...