Question

我正在定义一个函数，并希望使用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（），但是什么都做不到。我该怎么做？

Answer 1

如何在Numba Vectorize签名中指定元组？

在numba.vectorize函数中，不能使用元组。这是因为vectorize将这些类型的数组的代码矢量化。

因此，使用float, float, tuple签名将创建一个函数，该函数期望两个包含浮点数的数组和一个包含元组的数组。问题在于，包含元组的数组没有dtype-如果您使用结构化数组而不是包含元组的数组，但我没有尝试过，则可以使用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））。

如果您只想对float数组进行矢量化

如果您不想对元组参数的函数进行矢量化处理，则可以在另一个函数内部简单地创建矢量化函数，然后在其中调用它：

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"一起使用，您可能需要自己进行测试。

如何在Numba Vectorize签名中指定元组？

1 个答案: