如果 list_of_list[0] 中的元素不起作用，则以无 python 模式将列表列表传递给 numba 函数

Question

见下面的最小代码，

import numba
list_of_list = [[1, 2], [34, 100]]
@numba.njit()
def test(list_of_list):
    if 1 in list_of_list[0]:
        return 'haha'

test(list_of_list)

这行不通，似乎 list_of_list[0] 在编译期间不再表现得像列表。但是，以下代码有效：

list_of_list = [[1, 2], [34, 100]][0] # this is a list NOW!
@numba.njit()
def test(list_of_list):
    if 1 in list_of_list:
        return 'haha'

test(list_of_list)

这次我传入的实际上是列表，而不是列表的列表。然后就可以了。似乎 for i in list 在 numba 中有效，而不是 for i in list_of_list。

在我的用例中，将列表或数组的列表（如二维数据）传递给 numba 函数是很常见的。有时我只需要列表中的一个元素，这是在程序中动态确定的。

为了让它起作用，我其实想出了一个解决方案：把list_of_list压平成long list，然后用线性索引提取原始 list_of_list.

中的一个元素

我在这里问一下，还有其他替代方案吗？

Answer 1

in 方法适用于集合。返回一个字符串也会导致一些问题。

工作示例

import numba as nb
import numpy as np

array_2D = np.array([[1, 2], [34, 100]])

@nb.njit()
def test(array_2D):
    if 1 in set(array_2D[0]):
        #Strings also causes sometimes problems
        #return 'haha'
        return(1)
    else:
        return(-1)

Answer 2

你可以return一个带有我修改版本的字符串。它通过了测试并成功运行。

from numba import njit
import numpy as np


@njit
def test():
    if 1 in set(np_list_of_list[0]):
        return 'haha'


if __name__ == '__main__':
    list_of_list = [[1, 2], [34, 100]]
    np_list_of_list = np.array(list_of_list)
    print(test())

如果 list_of_list[0] 中的元素不起作用，则以无 python 模式将列表列表传递给 numba 函数

pass list of list into numba function in no python mode, if element in list_of_list[0] doe not work

numpy

python-3.x

numba

numba-pro