来自 numba 的 jit 和来自 scipy 的 hypergeom 的断言错误

Question

我正在尝试使用此功能：

@jit
def pv (matrix1, matrix2, vec, n):
    for i in range (n):
        for j in range (n):
            matrix1[i,j] = 1 - sum (hypergeom.pmf(X, n, int(vec[i]), int(vec[j])) for X in range(matrix2[i,j]))

但是我得到了 断言错误，很多事情我什至都没有开始理解（从 numba 回溯），最后 AssertionError：在对象模式管道中失败（步骤：对象模式前端）。我怀疑与 hyergeom 有关，但我看不出我做错了什么

编辑：最后我没有找到一种方法来实现使用@jit 的代码，但我找到了执行此操作的函数 scipy.stats.hypergeom.cdf：

sum (hypergeom.pmf(X, n, int(vec[i]), int(vec[j])) for X in range(matrix2[i,j]))

from scipy.stats import hypergeom

hypergeom.cdf(m2[i,j], n, v[i], v[j])

虽然这个解决方案加快了代码速度，但 for 循环仍然很慢（n =5053 需要半个多小时才能运行）

Answer 1

将 Numba 用于此类事情是个好主意，但不幸的是，正如您所怀疑的那样，它不支持 hypergeom 函数。您在这里没有做错任何事情 - 它只是不受支持，所以我认为在这种情况下您将无法使用 Numba。

支持的列表位于https://numba.pydata.org/numba-doc/latest/reference/numpysupported.html

我过去在发生这种情况时采取的一种方法是尝试使用 numba 确实支持但成功率可变的 numpy 子集编写我自己版本的不受支持的函数它可能会导致一系列全新的问题（为您自己的某些东西的实现换出经过调试、测试的库函数可能会导致垃圾箱火灾）。不看hypergeom.pmf的来源，我不知道这是否可行。