计算多维数组子矩阵之间的矩阵向量积

Question

我有一个形状为

A

的数组

(1000, 54, 50)

和另一个形状为

x

的数组

(1000, 54)

。对于每个

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))

Answer 1

最主要的是 matmul 将 N 维矩阵视为 2D 矩阵的堆栈。由于 x 已经是二维的，我们需要在最后广播一个额外的维度，将其视为一堆 54x1“矩阵”。最后的重塑是为了压平该维度。

import numpy as np

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

result = A @ A.transpose((0, 2, 1)) @ x[:, :, None]

result = result.reshape((A.shape[0:2]))