使用numba vectorize实现广播矩阵运算时报错：数组无法表示为Numpy类型

• nb.vectorize 仅支持输入标量对应输出标量的运算，你的函数返回1维数组自然会触发类型不支持的报错
• 之前写的guvectorize存在两个问题：一是把固定绑定的mat当成了动态输入参数，二是签名逻辑有误，且内部用@计算一维向量点积时新建了临时数组，才会触发连续性警告和结果错误

你需要用闭包动态生成绑定了固定mat的guvectorize函数，完全适配你动态生成大量矩阵对应计算函数的需求，实现代码如下：

import numba as nb
import numpy as np

def generate_load_calc_func(mat: np.ndarray):
    out_dim = mat.shape[0]  # 每次计算输出的一维数组长度
    # 签名定义：4个标量输入，输出长度为out_dim的一维数组
    @nb.guvectorize([f"void(f8, f8, f8, f8, f8[:])"], f"(),(),(),()->({out_dim})", nopython=True)
    def _calc(a, b, c, d, res):
        # 直接手动计算点积，不新建临时数组，消除性能警告
        for i in range(out_dim):
            res[i] = mat[i,0]*a + mat[i,1]*b + mat[i,2]*c + mat[i,3]*d
    return _calc

使用示例 #

# 生成绑定指定mat的计算函数
mat = np.array([[1.2, 1.6, 0.5, 0  ],
                [1.2, 1.6, 0,   0.5]])
f = generate_load_calc_func(mat)

# 标量输入测试
print(f(1.1, 2.1, 3, 4))
# 输出：[6.18 6.68] 符合预期

# 广播输入测试
a = np.array([1, 2, 3])
b = np.array([[1, 2, 3], [4, 5, 6]])
c = 1
d = 2
res = f(a, b, c, d)
print(res.shape)
# 输出：(2, 3, 2) 和你第一种预期的形状完全匹配
print(res)
"""
输出：
[[[ 3.3  3.8]
  [ 6.1  6.6]
  [ 8.9  9.4]]

 [[ 8.1  8.6]
  [10.9 11.4]
  [13.7 14.2]]]
"""

• 如果你需要第二种预期的输出形状，直接对结果做轴转置即可：res.transpose(2,0,1)就能得到你给出的第二种输出结构
• 如果mat的列数不固定，也可以调整实现逻辑，把4个输入参数合并为最后一维长度等于mat列数的数组，适配更通用的场景

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