多维矩阵中向量的点积（python，numpy）

Question

我有两个矩阵A，B，NxKxD维度，我想得到矩阵C，NxKxDxD维度，其中C [n，k] = A [n，k] x B [n，k] .T（此处“ x”表示乘积尺寸为Dx1和1xD的矩阵，因此结果必须为DxD维度），所以现在我的代码看起来像这样（这里A = B = X）：

def square(X):
    out = np.zeros((N, K, D, D))
    for n in range(N):
        for k in range(K):
            out[n, k] = np.dot(X[n, k, :, np.newaxis], X[n, k, np.newaxis, :])
    return out

对于大的N和K可能会比较慢，因为python的for周期。 有某种方法可以在一个numpy函数中进行这种乘法吗？

Answer 1

看来您不是在使用np.dot来减少总和，而只是为了进行扩展而导致广播。 因此，您可以简单地使用np.newaxis / None将数组扩展为一维，并让隐式广播提供帮助。

因此，一个实现将是-

X[...,None]*X[...,None,:]

关于广播的更多信息，特别是如何添加新轴，可以在this other post找到。