返回与传递的参数大小不同的数组的正确 numba 装饰器是什么？

Question

我想使用 numba 向量化一个函数来计算唯一值的出现次数。 该函数接收任意长度的 numpy 数组，并返回长度为 257 的 numpy 数组。

但是我不明白如何指定装饰器。它甚至无法编译。 我阅读了文档here，但我两手空空。它没有说明不同的数组大小。

@nb.guvectorize([(nb.uint8[:], nb.uint64[:])], "(n) -> (m)", target="parallel")
def count_occurrences(byte_view):
    """
    Counts the occurrences of each element in a byte array and returns a new array with the counts.
    """
    # adds a zero to the beginning of the array, for convenience
    count = np.zeros(1 + 256, dtype=np.uint64)
    count[1 + byte_view] += 1
    return count


sample = np.random.randint(1, 100, 100, dtype=np.uint8)
counts = count_occurrences(sample)

Answer 1

如果两者都是一维数组，则无法与此设置并行执行任何操作。在这种情况下，使用

guvectorize

似乎并没有比

njit

有多大好处。特别是因为

njit

允许返回形状与任何输入不同的数组。

不可能使用

guvectorize

返回未知形状的数组，因此您还必须将其作为输入参数提供。并且一开始就不应该从

guvectorize

函数显式返回任何内容。如果您认为 Numba（理想情况下）会并行执行您提供的函数，这可能会更直观，这意味着在这种情况下

count

将被初始化多次，而不是仅一次（预先）。

目前尚不清楚

byte_view

的索引在这种情况下意味着什么，如果它包含需要索引的索引，您可以如下所示循环它们。

import numpy as np
import numba as nb

@nb.guvectorize("void(uint8[:], uint64[:])", "(n),(m)", target="cpu")
def count_occurrences(byte_view, count):
    """
    Counts the occurrences of each element in a byte array and returns a new array with the counts.
    """
    for idx in byte_view: 
        count[1 + idx] += 1


sample = np.random.randint(1, 100, 100, dtype=np.uint8)

# adds a zero to the beginning of the array, for convenience
counts = np.zeros(1 + 256, dtype=np.uint64)

count_occurrences(sample, counts)

上面的代码片段使它可以工作，但正如所说，它并没有在

njit

上添加太多，并且没有利用

guvectorize

通常在

njit

上添加的内容。