如果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。

在我的用例中，将列表或数组如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())