numpy 如何并行使用numba来处理可变长度的列表

mgdq6dx1  于 2023-08-05  发布在  其他
关注(0)|答案(1)|浏览(134)

我有一个函数,它已经被njited并返回一个长度不同的uint32列表。

@numba.njit()
def inner_function(A: numba.float32[:]) -> List[int]:
    # ...

字符串
该函数在循环中被调用了大量次,这可以被并行化,因为没有数据依赖性。
但由于函数返回的长度不同,因此不能使用预分配数组。
我目前使用的是一种顺序方法,使用了像这样的numbas类型列表

@numba.njit()
def looping_function(A: numba.float32[:,:]) -> List[List[int]]:
    result = List()
    for i in range(A.shape[0]):
        tmp = inner_function(A[i])
        result.append(tmp)
    return result


有没有一种方法可以初始化我的结果列表(或任何其他数据结构),这样我就可以索引它,而不是像这样附加到它:

@numba.njit(parallel=True)
def looping_function(A: numba.float32[:,:]):
    result = ???
    for i in numba.prange(A.shape[0]):
        tmp = inner_function(A[i])
        result[i] = tmp
    return result

olmpazwi

olmpazwi1#

我找到了一个解决方案,它不是很令人满意,但对我的用例足够好。这将预分配为带有占位符0的列表,以确保类型推断。然后我可以删除占位符并根据需要进行扩展。

@numba.njit(parallel=True)#,forceobj=True)
def outer(i):
    a = numba.typed.List([numba.typed.List([0]) for _ in range(i)])
    for i in numba.prange(i):
        a[i].pop()
        a[i].extend(inner(i))
    return a

字符串

相关问题