如果list_of_list[0]中的元素不起作用，则在无python模式下将list的list传递到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。

在我的用例中，将列表或类似数组的2d数据列表传递到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

您可以使用我的修订版本返回一个字符串。它通过了测试并成功工作。

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())