如何使用Numba装饰器实现输入输出数组尺寸不同的函数？

你的代码无法编译的核心问题是错误使用了guvectorize装饰器，它的设计逻辑和你写的函数不匹配，下面是具体的问题分析和修复方案：

问题根源 #

1. guvectorize要求函数直接修改传入的输出数组，而不是创建数组后返回——你原来的return count完全不符合它的使用规范。
2. 签名里的(n) -> (m)无效，因为输出数组长度是固定的257，必须明确指定维度，不能用变量m。
3. target="parallel"在这里没必要，你是对单个输入数组做全局统计，并行化反而会增加额外开销。

修复方案 #

方案1：正确使用guvectorize #

按照guvectorize的规范，函数要接收输入数组和预先分配好的输出数组，直接填充输出数组：

import numba as nb
import numpy as np

@nb.guvectorize([(nb.uint8[:], nb.uint64[:])], "(n) -> (257)", target="cpu")
def count_occurrences(byte_view, counts):
    """统计字节数组中每个元素的出现次数，返回长度为257的计数数组（索引0预留）"""
    # 先重置计数数组为0
    counts[:] = 0
    # 遍历输入统计次数
    for b in byte_view:
        counts[1 + b] += 1

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

方案2：改用nb.jit更简单 #

如果你不需要向量化批量处理多个数组，用nb.jit直接编译普通函数更符合你的原始写法：

import numba as nb
import numpy as np

@nb.jit(nb.uint64[:](nb.uint8[:]), nopython=True)
def count_occurrences(byte_view):
    """统计字节数组中每个元素的出现次数，返回长度为257的计数数组（索引0预留）"""
    count = np.zeros(257, dtype=np.uint64)
    for b in byte_view:
        count[1 + b] += 1
    return count

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

额外优化 #

如果想让统计更快，可以用Numba支持的np.bincount替代循环，比如在guvectorize版本里修改：

@nb.guvectorize([(nb.uint8[:], nb.uint64[:])], "(n) -> (257)", target="cpu")
def count_occurrences(byte_view, counts):
    counts[:] = 0
    # 用bincount一次性统计，再放到对应位置
    bc = np.bincount(byte_view, minlength=256)
    counts[1:257] = bc

内容的提问来源于stack exchange，提问作者tutizeri