如何用np.vectorize或numpy数组广播消除for循环？

原代码示例 #

import numpy as np

xs = np.random.uniform(-2,2,size=(3,2))

Qs = np.array([
    [
        [4,2],
        [2,2]
    ],
    [
        [1,1],
        [1,1]
    ],
    [
        [2,2],
        [2,2]
    ]
])
rs = np.array([
    [1,-1],
    [1,1],
    [-1,1]
])

g = lambda Q, r, x: 0.5*x.T @ Q @ x + r @ x

# 原循环实现
results = []
for Q, r, x in zip(Qs, rs, xs):
    results.append(g(Q, r, x))
print(np.array(results))

现有尝试的错误原因 #

1. 直接广播调用失败：
   执行results = g(Qs, rs, xs)时，xs维度为(3,2)，x.T会转为(2,3)，而Qs是(3,2,2)。x.T @ Qs尝试执行(2,3) @ (3,2,2)，不符合矩阵乘法的维度匹配规则（matmul要求最后两个维度满足(n,k) @ (k,m) -> (n,m)，或前面维度可广播），因此报错维度不匹配。

2. np.vectorize调用失败：
   默认的np.vectorize会逐元素传递输入数组的每个元素，而非按样本维度传递完整的Q、r、x数组。这导致g函数接收到单个数值（如numpy.float64和numpy.int32），无法执行矩阵乘法@，因此抛出类型错误。

解决方案1：使用数组广播+维度调整（推荐，效率更高） #

通过调整数组维度，让numpy的广播机制支持批量矩阵运算，无需显式循环：

# 重写g为批量计算版本
def g_batch(Qs, rs, xs):
    # 将xs从(3,2)转为(3,2,1)，适配矩阵乘法的广播规则
    xs_col = xs[:, :, np.newaxis]
    # 计算二次项：0.5 * x.T @ Q @ x，结果从(3,1,1)压缩为(3,)
    quadratic = 0.5 * np.squeeze(xs_col.transpose(0,2,1) @ Qs @ xs_col)
    # 计算线性项：r与x的逐行点积，结果为(3,)
    linear = np.sum(rs * xs, axis=1)
    # 也可改用matmul实现：np.squeeze(rs @ xs_col)
    return quadratic + linear

# 调用批量版本
results_batch = g_batch(Qs, rs, xs)
print(results_batch)

解决方案2：正确使用np.vectorize（指定signature） #

通过signature参数告诉np.vectorize按维度分组传递参数，而非逐元素处理：

# 指定signature：每个Q是(2,2)矩阵，r和x是(2,)向量，输出为标量
g_vec = np.vectorize(g, signature='(n,n),(n),(n)->()')
results_vec = g_vec(Qs, rs, xs)
print(results_vec)

验证结果 #

两种方法的输出与原循环结果完全一致：

print(np.allclose(results, results_batch))  # True
print(np.allclose(results, results_vec))    # True

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