Numba cuda.jit多类型签名声明触发TypeError报错求助

错误原因 #

你遇到的TypeError是因为Numba 0.54版本的CUDA JIT装饰器要求多类型签名必须明确指定返回类型。你的签名仅包含参数类型，未声明无返回值对应的void类型，导致Numba无法解析签名格式，抛出“不是可调用对象”的错误。

修正方案 #

导入void类型，为每个签名添加返回类型声明，具体代码如下：

from numba import cuda 
from numba import float32, int32, void


@cuda.jit([(void, int32[:], int32[:], int32[:], int32),
           (void, float32[:], float32[:], int32[:], int32)])
def reset_when_done(data, ref, done, force_reset):
    env_id = cuda.blockIdx.x
    tid = cuda.threadIdx.x
    if tid == 0:
        if force_reset > 0.5 or done[env_id] > 0.5:
            data[env_id] = ref[env_id]

补充说明 #

• void表示函数无返回值，是Numba中对应无返回函数的类型标识；
• 每个签名元组的第一个元素为返回类型，后续元素依次对应函数参数的类型，需与函数定义的参数顺序严格匹配；
• 也可以通过signatures关键字参数传递多签名，效果完全一致：

  @cuda.jit(signatures=[(void, int32[:], int32[:], int32[:], int32),
                        (void, float32[:], float32[:], int32[:], int32)])
  def reset_when_done(data, ref, done, force_reset):
      # 函数体不变
  

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