在返回行 function 上获取 numba.core.errors.TypingError，尝试使用 numba.cuda 编译它

Question

我正在尝试用 Numba Cuda 制作 Erathostenes 筛子，以便用我的 GPU 完成大部分工作。 通过到处窃听一些代码，我得到了以下代码：

import datetime
import timeit

import numba
import numpy as np
from numba import *
from numba import cuda
import math


@cuda.jit(fastmath=True)
def sieve(primes: np.ndarray) -> 'np.ndarray[np.bool]':
    primes[0], primes[1] = False, False
    for i in numba.prange(2, int(math.sqrt(len(primes)) + 1)):
        if primes[i]:
            primes[i * i::i] = False
    return primes


if __name__ == "__main__":
    how_much = 1000000000
    threadsperblock = 32
    blockspergrid = (how_much + (threadsperblock - 1)) // threadsperblock
    timestart = timeit.default_timer()
    arr = np.ones(how_much, dtype=bool)
    print(np.flatnonzero(sieve[blockspergrid, threadsperblock](arr)))
    timestop = timeit.default_timer()
    timedelta = (timestop - timestart)
    print(f"Time Elapsed: {datetime.timedelta(seconds=timedelta)}")

但是我收到一个错误，好像 numba 无法编译 return 语句：

numba.core.errors.TypingError: Failed in cuda mode pipeline (step: nopython frontend)
No conversion from array(bool, 1d, C) to none for '$82return_value.1', defined at None

File "main.py", line 17:
def sieve(primes: np.ndarray) -> 'np.ndarray[np.bool]':
    <source elided>
            primes[i * i::i] = False
    return primes
    ^

During: typing of assignment at C:\Users\louis\PycharmProjects\RASCracker\main.py (17)

File "main.py", line 17:
def sieve(primes: np.ndarray) -> 'np.ndarray[np.bool]':
    <source elided>
            primes[i * i::i] = False
    return primes
    ^

我试图更改 function 的类型，但没有任何内容，使用 'np.ndarray[np.bool]' 或使用 np.ndarray 它不起作用。

是否有人知道如何修复此错误以正确编译筛子 function？