如何利用NumPy向量化高效计算多维数组子矩阵的矩阵-向量乘积？

算力阿峰

2026-7-1

问题：NumPy高效向量化实现矩阵运算

我有一个形状为(1000, 54, 50)的数组A，以及一个形状为(1000, 54)的数组x。需要对每个i=0,...,999，计算A[i, :, :] @ (A[i, :,:].T @ x[i])。请问如何借助NumPy的向量化能力实现最快的计算方式？

原有的循环实现（慢方法）：

import numpy as np

A = np.random.randn(1000, 54, 50)
x = np.random.randn(1000, 54)

def slow_method(A, x):
    B = np.zeros((1000, 54))
    for i in range(1000):
        B[i] = A[i] @ (A[i].T @ x[i])
    return B

我尝试过用einsum实现，但感觉存在大量冗余计算：

def einsum_method(A, x):
    return np.einsum('ijk,ik->ij', A, np.einsum('ijk,ik->ij', np.transpose(A, axes=(0, 2, 1)), x))

最优向量化实现

可以通过广播+矩阵乘法直接实现无冗余的向量化计算，彻底避免循环和重复运算：

def fast_vectorized(A, x):
    # 把x扩展为(1000,54,1)，适配矩阵乘法的维度要求
    x_reshaped = x[..., np.newaxis]
    # 批量计算每个样本的A[i].T @ x[i]，得到(1000,50,1)的中间结果
    intermediate = np.matmul(A.transpose(0, 2, 1), x_reshaped)
    # 批量计算A[i] @ intermediate，最后去掉多余的单维度
    return np.matmul(A, intermediate)[..., 0]

也可以写成更简洁的一行版：

def fast_vectorized_short(A, x):
    return A @ (A.transpose(0,2,1) @ x[..., None])[..., 0]

原理说明

维度适配：将x从(1000,54)扩展为(1000,54,1)，这样转置后的A（形状(1000,50,54)）能和它在batch维度（第0维）上自动广播，完成每个样本的A[i].T @ x[i]计算，得到(1000,50,1)的中间数组。
最终计算：用原始A（(1000,54,50)）和中间结果做矩阵乘法，得到(1000,54,1)，最后通过[...,0]去掉最后一个单维度，得到目标形状(1000,54)。