如何正确编写Numba @vectorize装饰器的多参数签名？

核心问题分析 #

你当前的代码存在几个关键误区：

1. n参数冗余：函数内部并未用到数组长度n，它只是生成输入数组的辅助变量，不需要作为@vectorize函数的参数——vectorize处理的是数组的单个元素，元素级操作和总长度无关。
2. 函数不能作为vectorize参数：exp、normalize、weigh、activate这些都是函数/向量化函数，@vectorize的参数只能是标量或数组（对应元素级输入），无法直接传递函数对象。
3. 签名定义错误：签名需要对应元素的类型，而非整个数组或函数的类型。

修正方案1：整合逻辑到单个vectorize函数 #

直接将所有元素级操作整合到一个@vectorize函数中，避免传递函数参数：

from numba import vectorize, cuda
import numpy as np

@vectorize('float32(float32, float32)', target='cuda')
def create_hidden_layer(greyscale, weight):
    # 实现normalize逻辑
    normalized = greyscale / 255
    # 实现weigh逻辑
    weighted = normalized * weight
    # 实现activate逻辑（直接用np.exp，numba会编译为CUDA兼容操作）
    a = np.exp(weighted)
    b = np.exp(-weighted)
    return (a - b) / (a + b)

• 签名说明：float32(float32, float32)表示输入两个float32类型的标量（对应greyscales和weights数组的单个元素），返回float32类型的结果。

修正方案2：复用已有标量编译函数 #

如果想复用normalize、weigh、activate的逻辑，需将它们改为标量级的@jit函数（而非@vectorize），再在@vectorize函数内部调用：

from numba import jit, vectorize, cuda
import numpy as np

# 定义标量级编译函数
@jit('float32(float32)', target='cuda')
def normalize(greyscale):
    return greyscale / 255

@jit('float32(float32, float32)', target='cuda')
def weigh(value, weight):
    return value * weight

@jit('float32(float32)', target='cuda')
def activate(value):
    a = np.exp(value)
    b = np.exp(-value)
    return (a - b) / (a + b)

# 定义向量化函数，调用标量函数
@vectorize('float32(float32, float32)', target='cuda')
def create_hidden_layer(greyscale, weight):
    normalized = normalize(greyscale)
    weighted = weigh(normalized, weight)
    activated = activate(weighted)
    return activated

调用方式 #

无论用哪种方案，最终调用时直接传入数组即可：

n = 1000000
greyscales = np.floor(np.random.uniform(0, 255, n).astype(np.float32))
weights = np.random.normal(.5, .1, n).astype(np.float32)

# 生成结果数组
hidden_layer = create_hidden_layer(greyscales, weights)

内容的提问来源于stack exchange，提问作者kndrtt