python - numpy 点积和矩阵积

Question

我正在使用形状为 (N,)、(N,3) 和 (N,3,3) 的 numpy 数组，它们表示 3D 空间中的标量、向量和矩阵序列。我已经实现了逐点点积、矩阵乘法和矩阵/向量乘法，如下所示:

def dot_product(v, w):
    return np.einsum('ij, ij -> i', v, w)

def matrix_vector_product(M, v):
    return np.einsum('ijk, ik -> ij', M, v)

def matrix_matrix_product(A, B):
    return np.einsum('ijk, ikl -> ijl', A, B)

如您所见，由于缺乏更好的解决方案，我使用了 einsum。令我惊讶的是我无法使用 np.dot... 这似乎不适合这种需要。有没有更 numpythonic 的方式来实现这些功能？

特别是，如果这些函数也可以通过广播第一个缺失的轴来处理形状 (3,) 和 (3,3)，那就太好了。我想我需要省略号，但我不太明白如何获得结果。

Answer 1

这些操作不能被 reshape 为一般的 BLAS 调用，并且循环 BLAS 调用对于这种大小的数组来说会非常慢。因此，einsum 可能是此类操作的最佳选择。

您的函数可以用省略号概括如下:

def dot_product(v, w):
    return np.einsum('...j,...j->...', v, w)

def matrix_vector_product(M, v):
    return np.einsum('...jk,...k->...j', M, v)

def matrix_matrix_product(A, B):
    return np.einsum('...jk,...kl->...jl', A, B)