将不同大小的数组列表传递给 Numba 函数

Question

我有一个预先计算的不同大小数组列表，我想将它传递给 Numba 函数。

from numba import jit
import numpy as np

@jit(nopython=True)
def go_fast(a, b): 
    ...
    return output

a = np.arange(100).reshape(10, 10)
b=[np.arange(4),np.arange(9)]

（实际上，b的元素是更复杂的数组，但这只是一个例子）。我怎样才能做到这一点？我知道 numba 不喜欢列表。

一种方法是将 b 变成带填充的高维数组，但提取实际元素需要循环，这并不理想。有没有更好的方法？

Answer 1

似乎新版本支持列表列表：

https://numba.pydata.org/numba-doc/dev/reference/pysupported.html

另一种选择是类型列表和类型字典：

https://numba.pydata.org/numba-doc/dev/reference/pysupported.html#dict https://numba.pydata.org/numba-doc/dev/reference/pysupported.html#list

将不同大小的数组列表传递给 Numba 函数

Pass list of differently sized arrays to Numba function

numba

numpy-ndarray