使用 numba 并行生成任意长度的随机向量

Question

我想生成如下形式的随机向量

[ i if random.uniform(0,1) <= probs[i] for i prange(K) ]

对于K 长度的概率数组probs。每个结果向量都有介于 0 和 K 之间的元素。从概念上讲，这就像翻转K 特定硬币（具有特定正面的概率）并记录哪些硬币显示正面。

任意返回长度使得难以使用numba 中的任何自动并行化选项。例如，

from numba import prange, njit, int64, float64
import numpy as np

@njit([int64[:](float64[:], int64)])
def rand_coin(freqs,r):
    return np.arange(r)[np.random.uniform(0,1,size=r)<=freqs] 

@njit(parallel=True)
def rand_coins(freqs,n):
    r = freqs.shape[0]
    return [rand_coin(freqs,r) for i in range(n)] #                     **

r = 10; n =100
freqs = np.random.uniform(0,1,r)
rand_coins(freqs, n) 

连续工作很好，但如果将星号行中的range 替换为prange，则会产生double free or corruption 错误。

是否可以并行化在numba 中返回随机长度数组的函数？

Answer 1

prange 是一个 numba 函数。输入错误只是 numba 的一般错误，表示它在编译函数时遇到了问题。真正的问题是您正在尝试调用未声明的函数。您需要像这样使用prange 函数：

from numba import njit, int64, float64
import numba
import numpy as np

@njit([int64[:](float64[:], int64)])
def rand_coin(freqs,r):
    return np.arange(r)[np.random.uniform(0,1,size=r)<=freqs] 

@njit(parallel=True)
def rand_coins(freqs,n):
    r = freqs.shape[0]
    return [rand_coin(freqs,r) for i in numba.prange(n)] #                     **