Skip to main content
 首页 » 编程设计

python之numba jitted 函数中集合的正确签名是什么

2024年09月03日41grandyang

如果我理解正确,我可以通过添加签名来提高 numba 函数的性能。示例:

@njit(int32(int32, int32)) 
def f(x, y): 
    # A somewhat trivial example 
    return x + y 

现在我有了需要两组的函数。正确的签名是什么?

@njit(int32(set(int32), set(int32))) 
def f(set_1, set_2): 
    # A somewhat trivial example 
    return x 

我认为签名 (int32(set(int32), set(int32))) 可能是正确的,但什么也没发生。 print(numba.typeof(set_1)) 返回 reflected set(int32)

请您参考如下方法:

If I understand correctly I can increase the performance of a numba function by adding a signature.

这是错误的——或者只是部分正确。使用签名 numba 只是提前编译函数,而不是在第一次使用这些参数调用时。第一次通话后,两者都应该同样快。在某些情况下,函数在没有签名的情况下可能会稍微快一些(尤其是对于 numba 可以使用输入的数组对齐的数组)。

Now I have function which takes two sets. What is the correct signature?

包含整数的 Python 集的正确签名是:

numba.types.Set(numba.int64, reflected=True) 

因此,接受两组(并返回一组)的函数的签名将是:

import numba as nb 
 
reflected_int_set = nb.types.Set(nb.int64, reflected=True) 
 
@nb.njit(reflected_int_set(reflected_int_set, reflected_int_set)) 
def f(set_1, set_2): 
    return set_1 
 
>>> f({1,2,3}, {3,4,5}) 
{1, 2, 3} 

但由于它(很可能)不会提高性能,所以我根本不会为签名烦恼。

还有一点要注意:numba 会在内部将 Python 集转换为 numba 集,因此将 Python set 传递给 numba 函数或从 numba 返回 set函数到 Python 上下文将复制完整的集合。在大多数情况下,开销比 numba 提供的潜在加速要重要得多。

根据我的经验,numba 的 set 和 list 只有在严格限于 numba 函数时才有意义。因此,如果您将它们用作参数或将它们返回(给非 numba 函数/上下文),您必须衡量性能并检查是否真的获得了加速。