Numba @vectorize装饰器多输出函数的签名如何定义？

首先明确：Numba的@vectorize装饰器不支持直接返回多个结果的元组，它的设计目标是对标量或一维输入进行向量化，输出单个标量或数组。针对你的需求，分两种场景给出解决方案：

场景1：velocity为固定值（如你代码中的[1.0,2.0,3.0]） #

既然velocity不依赖输入f，完全不需要把它放在向量化函数里，单独定义即可。用@vectorize只处理rho的计算：

import numba
import numpy as np

# 签名说明：输入是float64一维数组，输出是float64标量
@numba.vectorize("float64(float64[:])")
def compute_rho(f):
    return np.sum(f)

# 调用示例
f = np.random.rand(10, 10, 10, 5)  # 匹配你提到的[:,:,:,:]形状
rho = compute_rho(f)
# 直接定义固定的velocity
velocity = np.array([1.0, 2.0, 3.0], dtype=np.float64)

这里的签名"float64(float64[:])"告诉Numba：每个向量化单元接收一个一维float64数组（即f的最后一维），返回一个float64标量（求和后的rho）。装饰器会自动处理前面的维度，最终rho的形状和f[:,:,:,0]一致。

场景2：velocity需要依赖输入f计算 #

如果后续需要根据f动态计算u/v/w，改用@guvectorize装饰器——它是Numba专门用于多输入/多输出的向量化工具，支持处理数组块的维度映射：

import numba
import numpy as np

# 类型签名：(输入一维数组f, 输出标量rho, 输出一维数组velocity)
# 维度签名：'(n)->(),(m)' 表示输入n维块，输出一个0维标量和一个m维数组
@numba.guvectorize([(numba.float64[:], numba.float64, numba.float64[:])],
                   '(n)->(),(m)')
def rhou(f, rho_out, velocity_out):
    rho_out[0] = np.sum(f)
    # 这里可以根据f计算u/v/w，示例保留你的固定值逻辑
    velocity_out[0] = 1.0
    velocity_out[1] = 2.0
    velocity_out[2] = 3.0

# 调用示例
f = np.random.rand(10, 10, 10, 5)
rho, velocity = rhou(f)
# rho形状为(10,10,10)，velocity形状为(10,10,10,3)

维度签名'(n)->(),(m)'的作用是告诉Numba：将输入f的最后一维作为n进行处理，输出的rho是标量（对应()），velocity是长度为m的一维数组（这里m=3）。装饰器会自动广播处理前面的维度，最终输出和输入的前三维匹配。

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